diff --git a/bindings/node/Cargo.lock b/bindings/node/Cargo.lock index 359db4c8b..d7766e2da 100644 --- a/bindings/node/Cargo.lock +++ b/bindings/node/Cargo.lock @@ -89,6 +89,14 @@ version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +[[package]] +name = "bitsplit" +version = "0.1.0" +dependencies = [ + "ahash", + "atomsplit", +] + [[package]] name = "bitvec" version = "1.1.1" @@ -1436,6 +1444,7 @@ version = "0.23.2-dev.0" dependencies = [ "ahash", "atomsplit", + "bitsplit", "compact_str", "daachorse", "dary_heap", diff --git a/bindings/node/index.d.ts b/bindings/node/index.d.ts index 285f1b8b9..868c3d236 100644 --- a/bindings/node/index.d.ts +++ b/bindings/node/index.d.ts @@ -53,6 +53,21 @@ export declare class Normalizer { normalizeString(sequence: string): string } +export declare class PipelineTokenizer { + static fromFile(path: string): PipelineTokenizer + /** + * `Uint32Array`, not `Vec`: a JS `Array` costs one napi value per token, which + * on token-dense input is 13× the encode itself (gpt2 chinese 31 vs 616 MB/s). + */ + encode(text: string, addSpecialTokens?: boolean | undefined | null): Uint32Array + /** + * Drops the two remaining per-call costs of `encode`: the JS string → UTF-8 copy + * (as fast as the tokenizer itself, so it halves throughput) and the fresh + * ArrayBuffer (388 ns of a 789 ns call). Returns how many ids were written. + */ + encodeBytesInto(text: Uint8Array, out: Uint32Array, addSpecialTokens?: boolean | undefined | null): number +} + /** PreTokenizers */ export declare class PreTokenizer { preTokenizeString(sequence: string): [string, [number, number]][] diff --git a/bindings/node/index.js b/bindings/node/index.js index 7db5b31e6..0c7bb4a05 100644 --- a/bindings/node/index.js +++ b/bindings/node/index.js @@ -743,6 +743,7 @@ module.exports.Encoding = nativeBinding.Encoding module.exports.JsEncoding = nativeBinding.JsEncoding module.exports.Model = nativeBinding.Model module.exports.Normalizer = nativeBinding.Normalizer +module.exports.PipelineTokenizer = nativeBinding.PipelineTokenizer module.exports.PreTokenizer = nativeBinding.PreTokenizer module.exports.Processor = nativeBinding.Processor module.exports.Tokenizer = nativeBinding.Tokenizer diff --git a/bindings/node/src/lib.rs b/bindings/node/src/lib.rs index 5eb65461a..19bb04249 100644 --- a/bindings/node/src/lib.rs +++ b/bindings/node/src/lib.rs @@ -7,6 +7,7 @@ pub mod decoders; pub mod encoding; pub mod models; pub mod normalizers; +pub mod pipeline; pub mod pre_tokenizers; pub mod processors; pub mod tasks; diff --git a/bindings/node/src/pipeline.rs b/bindings/node/src/pipeline.rs new file mode 100644 index 000000000..a4d7b0e95 --- /dev/null +++ b/bindings/node/src/pipeline.rs @@ -0,0 +1,62 @@ +extern crate tokenizers as tk; + +use napi::bindgen_prelude::*; +use napi_derive::napi; +use tk::tokenizer::pipeline::PipelineTokenizer as Pipeline; + +fn err(e: E) -> Error { + Error::from_reason(format!("{e}")) +} + +#[napi] +pub struct PipelineTokenizer(Pipeline); + +#[napi] +impl PipelineTokenizer { + #[napi(factory)] + pub fn from_file(path: String) -> Result { + let tok = tk::Tokenizer::from_file(path).map_err(err)?; + Ok(Self(Pipeline::try_from(&tok).map_err(err)?)) + } + + /// `Uint32Array`, not `Vec`: a JS `Array` costs one napi value per token, which + /// on token-dense input is 13× the encode itself (gpt2 chinese 31 vs 616 MB/s). + #[napi] + pub fn encode(&self, text: String, add_special_tokens: Option) -> Result { + let ids = self + .0 + .encode(&text, add_special_tokens.unwrap_or(true)) + .map_err(err)?; + Ok(Uint32Array::new(ids.iter().map(|t| t.id).collect())) + } + + /// Drops the two remaining per-call costs of `encode`: the JS string → UTF-8 copy + /// (as fast as the tokenizer itself, so it halves throughput) and the fresh + /// ArrayBuffer (388 ns of a 789 ns call). Returns how many ids were written. + #[napi] + pub fn encode_bytes_into( + &self, + text: &[u8], + mut out: Uint32Array, + add_special_tokens: Option, + ) -> Result { + let text = std::str::from_utf8(text).map_err(err)?; + let ids = self + .0 + .encode(text, add_special_tokens.unwrap_or(true)) + .map_err(err)?; + // SAFETY: JS is blocked for this synchronous call, so nothing else aliases `out`. + let dst = unsafe { out.as_mut() }; + if ids.len() > dst.len() { + return Err(err(format!( + "need {} ids, buffer holds {}", + ids.len(), + dst.len() + ))); + } + for (d, t) in dst.iter_mut().zip(&ids) { + *d = t.id; + } + Ok(ids.len() as u32) + } +}