init

wasm/Cargo.toml

@@ -0,0 +1,60 @@
+[workspace]
+resolver = "3"
+rust-version = "1.94"
+edition = "2024"
+members = ["inference-models", "model-ai-text"]
+
+[workspace.dependencies]
+## Workspace models
+model-ai-text.path = "./model-ai-text"
+model-ai-text.version = "^0.1"
+## Wasm
+wasm-bindgen.version = "=0.2.114"
+wasm-bindgen-futures.version = "^0.4"
+js-sys.version = "^0.3"
+web-sys.version = "^0.3"
+gloo.version = "^0.11"
+wasm-tracing.version = "^2"
+console_error_panic_hook.version = "^0.1"
+## Logging
+tracing.version = "^0.1"
+tracing.features = ["max_level_debug", "release_max_level_info"]
+## Seri
+serde.version = "^1"
+serde_json.version = "^1"
+ciborium.version = "^0.2"
+## Codegen
+thiserror.version = "^2"
+## Inf
+tokenizers.version = "^0.22"
+tokenizers.default-features = false
+tokenizers.features = ["unstable_wasm"]
+burn.git = "https://github.com/tracel-ai/burn.git"
+burn.rev = "3a0e604cb14e445cd5a3686fbe06d0e11c76f05a"
+burn.default-features = false
+burn-store.git = "https://github.com/tracel-ai/burn.git"
+burn-store.rev = "3a0e604cb14e445cd5a3686fbe06d0e11c76f05a"
+burn-store.default-features = false
+burn-onnx.git = "https://github.com/tracel-ai/burn-onnx.git"
+burn-onnx.rev = "17a0864e02c19669b143e4227038363bb5194589"
+burn-onnx.default-features = false
+## HF
+hf-hub.version = "^0.5"
+
+[workspace.dependencies.getrandom_v4]
+package = "getrandom"
+version = "^0.4"
+features = ["wasm_js"]
+
+[workspace.dependencies.getrandom_v3]
+package = "getrandom"
+version = '^0.3'
+features = ["wasm_js"]
+
+[profile.release]
+opt-level = 3
+lto = "fat"
+codegen-units = 1
+panic = "abort"
+debug = "none"
+strip = "symbols"

wasm/inference-models/.cargo/config.toml

@@ -0,0 +1,7 @@
+[target.wasm32-unknown-unknown]
+rustflags = [
+    "-Ctarget-feature=+simd128,+bulk-memory,+reference-types",
+    '--cfg', 'getrandom_backend="wasm_js"',
+    "--cfg=has_std",
+    "-Zthreads=6"
+]

wasm/inference-models/Cargo.toml

@@ -0,0 +1,55 @@
+[package]
+name = "inference-models"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+## Models
+model-ai-text.workspace = true
+model-ai-text.optional = true
+## Wasm
+wasm-bindgen.workspace = true
+wasm-bindgen-futures.workspace = true
+js-sys.workspace = true
+gloo.workspace = true
+wasm-tracing.workspace = true
+console_error_panic_hook.workspace = true
+## Logging
+tracing.workspace = true
+## Inf
+burn.workspace = true
+burn.features = ["ndarray", "webgpu"]
+burn-store.workspace = true
+burn-store.features = ["burnpack"]
+tokenizers.workspace = true
+## Seri
+serde.workspace = true
+serde_json.workspace = true
+## Codegen
+thiserror.workspace = true
+## Wrkrnd
+getrandom_v4.workspace = true
+getrandom_v3.workspace = true
+
+[build-dependencies]
+hf-hub.workspace = true
+burn-onnx.workspace = true
+serde_json.workspace = true
+
+[features]
+ai_text = ["dep:model-ai-text"]
+
+[package.metadata.wasm-pack]
+profile.release.wasm-opt = [
+    "-O4",
+    "-ffm",
+    "--enable-simd",
+    "--enable-bulk-memory",
+    "--enable-reference-types",
+    "--enable-nontrapping-float-to-int",
+    "--strip-debug",
+    "--strip-producers",
+]

wasm/inference-models/src/capability.rs

@@ -0,0 +1,52 @@
+use std::sync::OnceLock;
+
+use burn::backend::{ndarray::NdArrayDevice, wgpu::WgpuDevice};
+use wasm_bindgen::prelude::*;
+
+/// This lock shall contain the prototype Device.
+/// `assert_*` functions shall set it once
+///
+/// If not set when check - assume default (`CubeDevice` implements Default as safe option)
+pub static CUBE_DEVICE: OnceLock<CubeDevice> = OnceLock::new();
+
+#[derive(Debug)]
+pub enum CubeDevice {
+    NdArray(NdArrayDevice),
+    WebGpu(WgpuDevice),
+}
+
+impl Default for CubeDevice {
+    fn default() -> Self {
+        Self::NdArray(Default::default())
+    }
+}
+
+#[wasm_bindgen(js_namespace = Capability)]
+pub async fn assert_webgpu() {
+    if let Some(d) = CUBE_DEVICE.get() {
+        tracing::warn!(asserted = ?d, "CUBE_DEVICE already set. Refuse to assert webgpu.");
+        return;
+    }
+
+    tracing::info!("Assert WebGPU-capable environment");
+    let device = Default::default();
+    burn::backend::wgpu::init_setup_async::<burn::backend::wgpu::graphics::WebGpu>(
+        &device,
+        Default::default(),
+    )
+    .await;
+
+    _ = CUBE_DEVICE.set(CubeDevice::WebGpu(device)).inspect_err(|e| tracing::error!(error=?e, "Burn backend was set to WGPU, but we can not testify this with CUBE_DEVICE"));
+}
+
+#[wasm_bindgen(js_namespace = Capability)]
+pub async fn assert_ndarray() {
+    if let Some(d) = CUBE_DEVICE.get() {
+        tracing::warn!(asserted = ?d, "CUBE_DEVICE already set. Refuse to assert ndarray.");
+        return;
+    }
+
+    tracing::info!("Assert NdArray environment");
+    let device = Default::default();
+    _ = CUBE_DEVICE.set(CubeDevice::NdArray(device)).inspect_err(|e| tracing::error!(error=?e, "Burn backend was set to NdArray, but we can not testify this with CUBE_DEVICE"));
+}

wasm/inference-models/src/inf/ai_text.rs

@@ -0,0 +1,259 @@
+use std::{
+    collections::VecDeque,
+    sync::{Arc, LazyLock, Mutex},
+};
+
+use burn::backend::{NdArray, WebGpu, ndarray::NdArrayDevice};
+use model_ai_text::*;
+use wasm_bindgen::prelude::*;
+
+use crate::capability::{CUBE_DEVICE, CubeDevice};
+
+static NDARRAY_FALLBACK_DEVICE: NdArrayDevice = NdArrayDevice::Cpu;
+
+#[allow(clippy::large_enum_variant)] // only used for runtime backend dispatch
+enum ModelVariant {
+    NdArray(Model<NdArray>),
+    WebGpu(Model<WebGpu>),
+}
+
+impl ModelVariant {
+    fn new_with_cube_device() -> Result<Self, ModelErr> {
+        let device = CUBE_DEVICE.get();
+        tracing::debug!("CUBE_DEVICE is {device:?} at the point of model creation");
+
+        match device {
+            Some(CubeDevice::NdArray(d)) => Model::try_load(d).map(Self::NdArray),
+            Some(CubeDevice::WebGpu(d)) => Model::try_load(d).map(Self::WebGpu),
+            None => {
+                tracing::warn!("No CUBE_DEVICE was asserted. Proceed with hope for NdArray");
+                Model::try_load(&NDARRAY_FALLBACK_DEVICE).map(Self::NdArray)
+            }
+        }
+    }
+
+    fn batch_inference(&self, texts: Vec<&str>) -> Result<Vec<ModelInference>, ModelInferenceErr> {
+        match self {
+            Self::NdArray(m) => m.batch_inference(texts),
+            Self::WebGpu(m) => m.batch_inference(texts),
+        }
+    }
+
+    fn stats(&self) -> ModelStats {
+        match self {
+            Self::NdArray(m) => m.stats(),
+            Self::WebGpu(m) => m.stats(),
+        }
+    }
+}
+
+type ModelMut = Arc<Mutex<Result<ModelVariant, ModelErr>>>;
+
+static MODEL: LazyLock<ModelMut> = LazyLock::new(|| {
+    tracing::debug!("Initializing AiText model with assumed CUBE_DEVICE");
+    Arc::new(Mutex::new(ModelVariant::new_with_cube_device()))
+});
+
+static INFERENCE_QUEUE: LazyLock<Arc<Mutex<VecDeque<AiTextInferenceTask>>>> =
+    LazyLock::new(|| Arc::new(Mutex::new(VecDeque::with_capacity(INFERENCE_QUEUE_SIZE))));
+
+static INFERENCE_QUEUE_SIZE: usize = 128;
+
+static INFERENCE_BATCH_SIZE: usize = 16;
+
+/// The AiText is merely a handle:
+/// - Shall pass the tasks into queue
+/// - Shall await for task execution on demand
+/// - Shall configure the model runtime (ndarray/webgpu)
+///
+/// This is pretty much a wasm-bindgen struct so expect not-very-rusty function signatures.
+#[wasm_bindgen]
+pub struct AiText;
+
+#[wasm_bindgen]
+impl AiText {
+    /// The handle shall force model load, signalizing the err
+    #[wasm_bindgen(constructor)]
+    #[allow(clippy::new_without_default)] // FIXME: introduce fail/recovery initialization
+    pub fn new() -> Self {
+        Self
+    }
+
+    /// Put the new inference task into the queue.
+    /// Future shall be resolved when task is moved out from queue and inference is done
+    ///
+    /// To get tasks from queue and run the inference call `run_queued_inference`
+    pub async fn queue(&self, text: &str) -> Result<AiTextInference, AiTextInferenceError> {
+        let task_future = {
+            let mut queue = INFERENCE_QUEUE.lock().expect("INFERENCE_QUEUE lock");
+            if queue.len() > INFERENCE_QUEUE_SIZE {
+                tracing::warn!("INFERENCE_QUEUE overflow: {INFERENCE_QUEUE_SIZE} limit reached");
+                return Err(AiTextInferenceError::ModelNotLoaded);
+            }
+
+            let task = AiTextInferenceTask {
+                text: text.into(),
+                async_bridge: Default::default(),
+            };
+            let future = task.future();
+
+            queue.push_back(task);
+
+            future
+        };
+
+        task_future.await
+    }
+
+    pub fn stats(&self) -> AiTextStats {
+        let model_stats = MODEL
+            .lock()
+            .expect("MODEL lock")
+            .as_ref()
+            .inspect_err(
+                |e| tracing::error!(error=?e, "Model didn't load: model stats shall be defaults"),
+            )
+            .map(|m| m.stats())
+            .ok();
+
+        let tokenizer_max_length = model_stats
+            .as_ref()
+            .and_then(|s| s.tokenizer_max_length)
+            .unwrap_or_default();
+
+        AiTextStats {
+            tokenizer_max_length,
+            queue_limit: INFERENCE_QUEUE_SIZE,
+            queue_size: INFERENCE_QUEUE.lock().expect("INFERENCE_QUEUE lock").len(),
+        }
+    }
+
+    /// Run inference for batch from the queue
+    pub async fn run_queued_inference(&self) {
+        let mut queue = INFERENCE_QUEUE.lock().expect("INFERENCE_QUEUE lock");
+
+        let model_lock = MODEL.lock().expect("MODEL lock");
+        let Ok(model) = model_lock
+            .as_ref()
+            .inspect_err(|e| tracing::error!(error=?e, "Can not run inference: model didn't load"))
+        else {
+            return;
+        };
+        let tasks = {
+            let size = queue.len().min(INFERENCE_BATCH_SIZE);
+            let mut tasks = Vec::with_capacity(size);
+
+            while tasks.len() < size {
+                if let Some(task) = queue.pop_front() {
+                    tasks.push(task)
+                } else {
+                    break;
+                }
+            }
+
+            tasks
+        };
+        let texts = Vec::from_iter(tasks.iter().map(|t| t.text.as_str()));
+
+        tracing::info!(size = texts.len(), "Request inference");
+        let result = model.batch_inference(texts);
+
+        match result {
+            Ok(probs) => {
+                tracing::info!(size = probs.len(), "Inference complete");
+                probs.into_iter().zip(tasks).for_each(|(prob, task)| {
+                    task.set(Ok(AiTextInference {
+                        prob_ai: prob.prob_ai,
+                        prob_human: prob.prob_human,
+                    }));
+                });
+            }
+            Err(e) => {
+                tracing::error!(error=?e, "Inference error");
+                tasks
+                    .into_iter()
+                    .for_each(|t| t.set(Err(AiTextInferenceError::InferenceExecution)));
+            }
+        }
+    }
+}
+
+#[wasm_bindgen]
+pub struct AiTextStats {
+    pub tokenizer_max_length: usize,
+    pub queue_limit: usize,
+    pub queue_size: usize,
+}
+
+#[wasm_bindgen]
+pub struct AiTextInference {
+    pub prob_ai: f32,
+    pub prob_human: f32,
+}
+
+#[wasm_bindgen]
+#[derive(Debug, Clone, thiserror::Error)]
+pub enum AiTextInferenceError {
+    #[error("Inference execution Error. Check logs for details")]
+    InferenceExecution,
+    #[error("Model is not loaded. Check logs for details")]
+    ModelNotLoaded,
+}
+
+/// This shall be the promise-like thingey, holded by the queue until
+/// explicitly disposed
+pub struct AiTextInferenceTask {
+    text: String,
+    /// This is the field containing inference result, which does also
+    /// act like a future output (promise resolve)
+    async_bridge: Arc<Mutex<InferenceTaskAsyncBridge>>,
+}
+
+impl AiTextInferenceTask {
+    fn future(&self) -> InferenceTaskPromisedFuture {
+        InferenceTaskPromisedFuture(self.async_bridge.clone())
+    }
+
+    fn set(&self, result: Result<AiTextInference, AiTextInferenceError>) {
+        let mut bridge = self.async_bridge.lock().expect("async_bridge lock");
+        bridge.result = Some(result);
+
+        if let Some(waker) = bridge.waker.take() {
+            waker.wake();
+        }
+    }
+}
+
+#[wasm_bindgen]
+#[derive(Default)]
+struct InferenceTaskAsyncBridge {
+    result: Option<Result<AiTextInference, AiTextInferenceError>>,
+    waker: Option<std::task::Waker>,
+}
+
+/// Future which shall resolve into promise-compatible thing, and trace the error (if any)
+/// coz error gets erased into JsValue
+struct InferenceTaskPromisedFuture(Arc<Mutex<InferenceTaskAsyncBridge>>);
+
+impl Future for InferenceTaskPromisedFuture {
+    type Output = Result<AiTextInference, AiTextInferenceError>;
+
+    fn poll(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        let mut inner = self.0.lock().expect("Inference future inner lock");
+
+        match inner.result.take() {
+            Some(result) => {
+                std::task::Poll::Ready(result.inspect_err(
+                    |e| tracing::error!(error=?e, "Failed to perform AiText inference"),
+                ))
+            }
+            None => {
+                inner.waker = Some(cx.waker().clone());
+                std::task::Poll::Pending
+            }
+        }
+    }
+}

wasm/inference-models/src/inf/mod.rs

@@ -0,0 +1,2 @@
+#[cfg(feature = "ai_text")]
+pub mod ai_text;

wasm/inference-models/src/lib.rs

@@ -0,0 +1,32 @@
+#![recursion_limit = "256"]
+
+use wasm_bindgen::prelude::*;
+
+pub mod capability;
+pub mod inf;
+
+#[wasm_bindgen]
+extern "C" {
+    #[wasm_bindgen(js_name = log, js_namespace = console)]
+    pub(crate) fn jsconsole_log(s: &str);
+    #[wasm_bindgen(js_name = error, js_namespace = console)]
+    pub(crate) fn jsconsole_err(s: &str);
+}
+
+#[wasm_bindgen(start)]
+pub fn start() {
+    console_error_panic_hook::set_once();
+    wasm_tracing::set_as_global_default();
+
+    let feature_report = std::format!(
+        "inference-models: ai-text {ai_text}",
+        ai_text = cfg_select! {
+            feature = "ai_text" => "✅",
+            _ => "❌"
+        }
+    );
+
+    tracing::info!(
+        "inference-models looks good;\nfeatures: {feature_report};\nThe extension source code is available at source.spushk.ru/pub/firaid-moz\nHappy browsing!"
+    );
+}

wasm/model-ai-text/Cargo.toml

@@ -0,0 +1,31 @@
+[package]
+name = "model-ai-text"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+## Logging
+tracing.workspace = true
+## Inf
+burn.workspace = true
+burn-store.workspace = true
+burn-store.features = ["burnpack"]
+tokenizers.workspace = true
+## Seri
+serde.workspace = true
+serde_json.workspace = true
+ciborium.workspace = true
+## Codegen
+thiserror.workspace = true
+
+[build-dependencies]
+hf-hub.workspace = true
+burn-onnx.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tokenizers.workspace = true
+ciborium.workspace = true
+
+[features]
+_has_attention = []
+_has_types = []

wasm/model-ai-text/build.rs

@@ -0,0 +1,203 @@
+use burn_onnx::ModelGen;
+use serde::{Deserialize, Deserializer};
+use std::{env, path::Path};
+use tokenizers::Tokenizer;
+
+#[derive(Deserialize)]
+pub struct Config {
+    hidden_size: usize,
+    num_attention_heads: usize,
+}
+
+#[derive(Deserialize)]
+pub struct TokenizerConfig {
+    pad_token: String,
+    #[serde(
+        alias = "model_max_length",
+        default = "TokenizerConfig::max_length_default",
+        deserialize_with = "TokenizerConfig::max_length_de"
+    )]
+    max_length: usize,
+}
+
+impl TokenizerConfig {
+    const fn max_length_default() -> usize {
+        512
+    }
+
+    /// If value is far above `LENGTH_CAP` - that means broken config. We're going to
+    /// play safe and replace it with `max_length_default`.
+    ///
+    /// Yeah yeah, I know, "(in)sane defaults".
+    ///
+    /// Yet our current model BTW (!) has insanely large `model_max_length` while
+    /// in reality its cap is 512.
+    fn max_length_de<'de, D>(deserializer: D) -> Result<usize, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct Visitor;
+        const LENGTH_CAP: usize = 8192;
+
+        impl<'de> serde::de::Visitor<'de> for Visitor {
+            type Value = usize;
+
+            fn expecting(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+                write!(
+                    fmt,
+                    "unsigned integer. Values are capped by {LENGTH_CAP}. Float is accepted, and also capped to {LENGTH_CAP}."
+                )
+            }
+
+            fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                // This is configuration sanity check.
+                if v > (LENGTH_CAP * 3) as u64 {
+                    println!(
+                        "WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
+                        TokenizerConfig::max_length_default()
+                    );
+                    Ok(TokenizerConfig::max_length_default())
+                } else {
+                    Ok(v.clamp(0, LENGTH_CAP as u64) as usize)
+                }
+            }
+
+            fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                // This is configuration sanity check.
+                if v > (LENGTH_CAP * 3) as i64 {
+                    println!(
+                        "WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
+                        TokenizerConfig::max_length_default()
+                    );
+                    Ok(TokenizerConfig::max_length_default())
+                } else {
+                    Ok(v.clamp(0, LENGTH_CAP as i64) as usize)
+                }
+            }
+
+            fn visit_f64<E>(self, v: f64) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                // This is configuration sanity check.
+                if v > (LENGTH_CAP * 3) as f64 {
+                    println!(
+                        "WARN: the parsed `max_length` {v} is well above {LENGTH_CAP}: falling back to {}",
+                        TokenizerConfig::max_length_default()
+                    );
+                    Ok(TokenizerConfig::max_length_default())
+                } else {
+                    Ok(v.clamp(0.0, LENGTH_CAP as f64) as usize)
+                }
+            }
+        }
+
+        deserializer.deserialize_any(Visitor)
+    }
+}
+
+fn main() {
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT");
+    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_ATTENTION");
+    println!("cargo:rerun-if-env-changed=_MK_MODEL_AI_TEXT_TYPES");
+
+    let model_repo = std::env::var("_MK_MODEL_AI_TEXT")
+        .unwrap_or_else(|_| "onnx-community/e5-small-lora-ai-generated-detector-ONNX".into());
+    let has_attention = std::env::var("_MK_MODEL_AI_TEXT_ATTENTION")
+        .is_ok_and(|v| v == "true" || v == "yes" || v == "1");
+    let has_types = std::env::var("_MK_MODEL_AI_TEXT_TYPES")
+        .is_ok_and(|v| v == "true" || v == "yes" || v == "1");
+
+    if has_attention {
+        println!("cargo:rustc-cfg=feature=\"_has_attention\"");
+    }
+    if has_types {
+        println!("cargo:rustc-cfg=feature=\"_has_types\"");
+    }
+
+    let api = hf_hub::api::sync::Api::new().expect("Huggingface API");
+
+    // dwnl model
+    let repo = api.model(model_repo);
+    let model = repo
+        .get("onnx/model.onnx")
+        .expect("ai_text: model.onnx fetch");
+    let config_json = repo.get("config.json").expect("ai_text: config.json fetch");
+    let tokenizer_json = repo
+        .get("tokenizer.json")
+        .expect("ai_text: tokenizer.json fetch");
+    let tokenizer_config_json = repo
+        .get("tokenizer_config.json")
+        .expect("ai_text: tokenizer_config.json fetch");
+
+    let burned_dir = Path::new("burned").join("ai_text");
+    let out_dir = Path::new(&env::var("OUT_DIR").expect("OUT_DIR is env var")).join(&burned_dir);
+
+    // handle model embedding
+    ModelGen::new()
+        .load_strategy(burn_onnx::LoadStrategy::Embedded)
+        .simplify(true)
+        .input(model.to_str().unwrap())
+        .out_dir(burned_dir.to_str().unwrap())
+        .run_from_script();
+
+    // construct tokenizer AoT
+    let tokenizer = {
+        let mut tokenizer = Tokenizer::from_file(tokenizer_json)
+            .inspect_err(|e| eprintln!("While loading tokenizer.json: {e:?}"))
+            .expect("Tokenizer JSON at tokenizer.json");
+
+        let tokenizer_config: TokenizerConfig = serde_json::from_reader(
+            std::fs::File::open(tokenizer_config_json)
+                .inspect_err(|e| eprintln!("While opening tokenizer_config.json: {e:?}"))
+                .expect("File tokenizer_config.json"),
+        )
+        .inspect_err(|e| eprint!("While deserializing tokenizer_config.json: {e:?}"))
+        .expect("Valid JSON at tokenizer_config.json");
+
+        let pad_id = tokenizer
+            .token_to_id(&tokenizer_config.pad_token)
+            .expect("Tokenizer to know PAD token");
+
+        let config: Config = serde_json::from_reader(
+            std::fs::File::open(config_json)
+                .inspect_err(|e| eprintln!("While opening config.json: {e:?}"))
+                .expect("Open file config.json"),
+        )
+        .inspect_err(|e| eprintln!("While deserializing config.json: {e:?}"))
+        .expect("Valid JSON at config.json");
+
+        tokenizer
+            .with_truncation(Some(tokenizers::TruncationParams {
+                direction: tokenizers::TruncationDirection::Right,
+                max_length: tokenizer_config.max_length,
+                strategy: Default::default(),
+                stride: Default::default(),
+            }))
+            .inspect_err(|e| eprintln!("While setting tokenizer truncation: {e:?}"))
+            .expect("Set tokenizer truncation");
+
+        tokenizer.with_padding(Some(tokenizers::PaddingParams {
+            strategy: tokenizers::PaddingStrategy::BatchLongest,
+            direction: tokenizers::PaddingDirection::Right,
+            pad_to_multiple_of: Some(config.hidden_size / config.num_attention_heads),
+            pad_id,
+            pad_type_id: 0,
+            pad_token: tokenizer_config.pad_token,
+        }));
+
+        tokenizer
+    };
+
+    tokenizer
+        .save(out_dir.join("tokenizer_d.json"), false)
+        .inspect_err(|e| eprintln!("While saving tokenizer: {e:?}"))
+        .expect("Save tokenizer");
+}

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,
+            },
+        )))
+    }
+}