init

wasm/model-ai-text/src/lib.rs

@@ -0,0 +1,148 @@
+#![recursion_limit = "256"]
+#![feature(iter_array_chunks)]
+
+use burn::{prelude::*, tensor::activation::quiet_softmax};
+use tokenizers::Tokenizer;
+
+mod data {
+    #![allow(clippy::all)] // have massive codegen here.
+    extern crate alloc;
+
+    include!(concat!(env!("OUT_DIR"), "/burned/ai_text/model.rs"));
+
+    pub static TOKENIZER: &[u8] =
+        include_bytes!(concat!(env!("OUT_DIR"), "/burned/ai_text/tokenizer_d.json"));
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ModelErr {
+    #[error("Failed to load TOKENIZER")]
+    Tokenizer(Box<dyn core::error::Error + Send + Sync>),
+    #[error("Invalid TOKENIZER_CONFIG")]
+    TokenizerConfig(serde_json::Error),
+    #[error("Invalid CONFIG")]
+    Config(serde_json::Error),
+    #[error("Missing required token `{0}` (this is name, not the token itself)")]
+    MissingToken(&'static str),
+    #[error("Failed to apply truncation to the tokenizer")]
+    TokenizerWithTruncation(Box<dyn core::error::Error + Send + Sync>),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ModelInferenceErr {
+    #[error("Failed to encode inputs")]
+    InputEncode(Box<dyn core::error::Error + Send + Sync>),
+}
+
+#[derive(Debug, Clone)]
+pub struct ModelInference {
+    pub prob_ai: f32,
+    pub prob_human: f32,
+}
+
+#[derive(Debug, Clone)]
+pub struct ModelStats {
+    pub tokenizer_max_length: Option<usize>,
+}
+
+pub struct Model<B: Backend> {
+    model: data::Model<B>,
+    tokenizer: Tokenizer,
+    device: &'static <B as Backend>::Device,
+}
+
+impl<B: Backend> Model<B> {
+    pub fn try_load(device: &'static <B as Backend>::Device) -> Result<Self, ModelErr> {
+        tracing::debug!(device = ?device, "Requested to load model and tokenizer");
+
+        let tokenizer = Tokenizer::from_bytes(data::TOKENIZER).map_err(ModelErr::Tokenizer)?;
+        let model = data::Model::from_embedded(device);
+
+        Ok(Self {
+            device,
+            model,
+            tokenizer,
+        })
+    }
+
+    pub fn stats(&self) -> ModelStats {
+        let tokenizer_max_length = self.tokenizer.get_truncation().map(|t| t.max_length);
+
+        ModelStats {
+            tokenizer_max_length,
+        }
+    }
+
+    #[tracing::instrument(skip_all)]
+    pub fn batch_inference(
+        &self,
+        texts: Vec<&str>,
+    ) -> Result<Vec<ModelInference>, ModelInferenceErr> {
+        tracing::debug!(len = texts.len(), "Requested batch inference");
+
+        if texts.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        let encodings = self
+            .tokenizer
+            .encode_batch(texts, true)
+            .map_err(ModelInferenceErr::InputEncode)?;
+
+        let batch_len = encodings.len();
+        let encoding_len = encodings.first().map(|f| f.len()).unwrap_or_default();
+
+        let (ids, attn, types) = {
+            let bulk_size = batch_len * encoding_len;
+
+            let mut buf = vec![0u32; bulk_size * 3]; // 3 bulks: ids, attn, types
+
+            let ids_i = 0;
+            let attn_i = ids_i + bulk_size;
+            let types_i = attn_i + bulk_size;
+
+            for (i, e) in encodings.iter().enumerate() {
+                let off = i * encoding_len;
+                buf[ids_i + off..ids_i + off + encoding_len].copy_from_slice(e.get_ids());
+                buf[attn_i + off..attn_i + off + encoding_len]
+                    .copy_from_slice(e.get_attention_mask());
+                buf[types_i + off..types_i + off + encoding_len].copy_from_slice(e.get_type_ids());
+            }
+
+            let shape = Shape::new([batch_len, encoding_len]);
+
+            let ids = &buf[ids_i..ids_i + bulk_size];
+            let attn = &buf[attn_i..attn_i + bulk_size];
+            let types = &buf[types_i..types_i + bulk_size];
+
+            let ids: Tensor<B, 2, Int> =
+                Tensor::<B, 1, Int>::from_ints(ids, self.device).reshape(shape.clone());
+
+            let attn: Tensor<B, 2, Int> =
+                Tensor::<B, 1, Int>::from_ints(attn, self.device).reshape(shape.clone());
+
+            let types: Tensor<B, 2, Int> =
+                Tensor::<B, 1, Int>::from_ints(types, self.device).reshape(shape);
+
+            (ids, attn, types)
+        };
+
+        tracing::debug!("Batch inference ready, encoded");
+        let logits = cfg_select! {
+            all(feature = "_has_attention", feature = "_has_types") => self.model.forward(ids, attn, types),
+            feature = "_has_types" => self.model.forward(ids, types),
+            feature = "_has_attention" => self.model.forward(ids, attn),
+            _ => self.model.forward(ids)
+        };
+
+        tracing::debug!(logits = ?logits, "Batch inference complete");
+        let probs = quiet_softmax(logits, 1).into_data();
+
+        Ok(Vec::from_iter(probs.iter::<f32>().array_chunks::<2>().map(
+            |[prob_human, prob_ai]| ModelInference {
+                prob_ai,
+                prob_human,
+            },
+        )))
+    }
+}