TARGET MINIMAL RUST CRATE + .tok FORMAT

bindings/node-tok/Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "node-tok"
+version = "0.23.2-dev.0"
+edition = "2024"
+authors = ["Arthur Zucker <arthur@huggingface.co>"]
+license = "Apache-2.0"
+description = "Node binding whose entire surface is the `.tok` read path, so no JSON parser links."
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+napi = { version = "3", default-features = false, features = ["napi6"] }
+napi-derive = "3"
+# No `config`: this binding loads `.tok` and nothing else, so the JSON layer and every
+# enum arm that only it can construct stay out of the binary.
+tk-encode = { path = "../../tokenizers/tk-encode", default-features = false }
+tk-serialization = { path = "../../tokenizers/tk-serialization" }
+
+[build-dependencies]
+napi-build = "2"
+
+[profile.release]
+lto = "fat"
+strip = true
+panic = "abort"
+codegen-units = 1
+opt-level = "z"
+
+# This package is outside the workspace, so it needs the patch spelled out again.
+# https://github.com/RagnarGrootKoerkamp/PtrHash/pull/32 — drop once it lands.
+[patch.crates-io]
+ptr_hash = { git = "https://github.com/ArthurZucker/PtrHash", branch = "feat/optional-rayon" }

bindings/node-tok/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    napi_build::setup();
+}

bindings/node-tok/src/lib.rs

@@ -0,0 +1,72 @@
+//! A Node binding whose entire surface is the `.tok` read path.
+//!
+//! There is deliberately no `fromFile`, no `fromString`, no config accessors: the moment one
+//! exists, `serde_json` is reachable and LTO has to keep the whole JSON stack. Converting a
+//! `tokenizer.json` is `tk-convert`'s job, offline.
+
+use napi::bindgen_prelude::*;
+use napi_derive::napi;
+
+use tk_encode::pipeline::PipelineTokenizer as Pipeline;
+
+#[napi]
+pub struct TokTokenizer {
+    inner: Pipeline,
+    // Keeps the mapped bytes alive for as long as the tokenizer that was built from them.
+    _file: tk_serialization::TokFile,
+}
+
+#[napi]
+impl TokTokenizer {
+    /// Load a `.tok` produced by `tk-convert`.
+    #[napi(factory)]
+    pub fn from_file(path: String) -> Result<Self> {
+        let file = tk_serialization::TokFile::open(&path)
+            .map_err(|e| Error::from_reason(format!("{path}: {e}")))?;
+        let inner = Pipeline::from_tok(file.bytes())
+            .map_err(|e| Error::from_reason(format!("{path}: {e}")))?;
+        Ok(Self { inner, _file: file })
+    }
+
+    /// Encode `text`, returning the ids. A `Uint32Array` rather than a `Vec<u32>`: the latter
+    /// marshals as a boxed JS array, one napi value per token, which costs more than the encode.
+    #[napi]
+    pub fn encode(&self, text: String, add_special_tokens: bool) -> Result<Uint32Array> {
+        let encoded = self
+            .inner
+            .encode(text.as_str(), add_special_tokens)
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+        Ok(encoded.iter().map(|t| t.id).collect::<Vec<u32>>().into())
+    }
+
+    /// Encode UTF-8 bytes straight into a caller-owned buffer, returning how many ids were
+    /// written. Drops the JS string copy and the fresh ArrayBuffer.
+    #[napi]
+    pub fn encode_bytes_into(
+        &self,
+        text: Buffer,
+        mut out: Uint32Array,
+        add_special_tokens: bool,
+    ) -> Result<u32> {
+        let text = std::str::from_utf8(&text)
+            .map_err(|e| Error::from_reason(format!("input is not valid UTF-8: {e}")))?;
+        let encoded = self
+            .inner
+            .encode(text, add_special_tokens)
+            .map_err(|e| Error::from_reason(e.to_string()))?;
+        if encoded.len() > out.len() {
+            return Err(Error::from_reason(format!(
+                "output buffer holds {} ids, needs {}",
+                out.len(),
+                encoded.len()
+            )));
+        }
+        // SAFETY: `out` is a JS-owned `Uint32Array` handed to this call; napi only marks the
+        // mutable view unsafe because JS could alias it, and nothing here re-enters JS.
+        let slots = unsafe { out.as_mut() };
+        for (slot, token) in slots.iter_mut().zip(&encoded) {
+            *slot = token.id;
+        }
+        Ok(encoded.len() as u32)
+    }
+}