init

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!"
+    );
+}