spnl 0.6.0__tar.gz → 0.6.1__tar.gz

{spnl-0.6.0 → spnl-0.6.1}/Cargo.lock

@@ -149,9 +149,9 @@ dependencies = [
 
 [[package]]
 name = "anyhow"
-version = "1.0.99"
+version = "1.0.100"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b0674a1ddeecb70197781e945de4b3b8ffb61fa939a5597bcf48503737663100"
+checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61"
 
 [[package]]
 name = "arc-swap"
@@ -2890,7 +2890,7 @@ dependencies = [
 
 [[package]]
 name = "haystack"
-version = "0.6.0"
+version = "0.6.1"
 dependencies = [
  "anyhow",
  "clap",
@@ -3500,9 +3500,9 @@ dependencies = [
 
 [[package]]
 name = "js-sys"
-version = "0.3.78"
+version = "0.3.80"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0c0b063578492ceec17683ef2f8c5e89121fbd0b172cbc280635ab7567db2738"
+checksum = "852f13bec5eba4ba9afbeb93fd7c13fe56147f055939ae21c43a29a0ecb2702e"
 dependencies = [
  "once_cell",
  "wasm-bindgen",
@@ -5031,7 +5031,7 @@ checksum = "9cd31dcfdbbd7431a807ef4df6edd6473228e94d5c805e8cf671227a21bad068"
 dependencies = [
  "anyhow",
  "clap",
- "itertools 0.14.0",
+ "itertools 0.13.0",
  "proc-macro2",
  "quote",
  "rand 0.8.5",
@@ -5218,7 +5218,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf"
 dependencies = [
  "heck",
- "itertools 0.14.0",
+ "itertools 0.13.0",
  "log",
  "multimap",
  "once_cell",
@@ -5238,7 +5238,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d"
 dependencies = [
  "anyhow",
- "itertools 0.14.0",
+ "itertools 0.13.0",
  "proc-macro2",
  "quote",
  "syn 2.0.106",
@@ -6158,10 +6158,11 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc"
 
 [[package]]
 name = "serde"
-version = "1.0.219"
+version = "1.0.225"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6"
+checksum = "fd6c24dee235d0da097043389623fb913daddf92c76e9f5a1db88607a0bcbd1d"
 dependencies = [
+ "serde_core",
  "serde_derive",
 ]
 
@@ -6175,11 +6176,20 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "serde_core"
+version = "1.0.225"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "659356f9a0cb1e529b24c01e43ad2bdf520ec4ceaf83047b83ddcc2251f96383"
+dependencies = [
+ "serde_derive",
+]
+
 [[package]]
 name = "serde_derive"
-version = "1.0.219"
+version = "1.0.225"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00"
+checksum = "0ea936adf78b1f766949a4977b91d2f5595825bd6ec079aa9543ad2685fc4516"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -6188,14 +6198,15 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.143"
+version = "1.0.145"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d401abef1d108fbd9cbaebc3e46611f4b1021f714a0597a71f41ee463f5f4a5a"
+checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c"
 dependencies = [
  "itoa",
  "memchr",
  "ryu",
  "serde",
+ "serde_core",
 ]
 
 [[package]]
@@ -6483,7 +6494,7 @@ dependencies = [
 
 [[package]]
 name = "spnl"
-version = "0.6.0"
+version = "0.6.1"
 dependencies = [
  "anyhow",
  "arrow-array",
@@ -6521,7 +6532,7 @@ dependencies = [
 
 [[package]]
 name = "spnl-cli"
-version = "0.6.0"
+version = "0.6.1"
 dependencies = [
  "anyhow",
  "clap",
@@ -6536,7 +6547,7 @@ dependencies = [
 
 [[package]]
 name = "spnl-wasm"
-version = "0.6.0"
+version = "0.6.1"
 dependencies = [
  "serde_json",
  "spnl",
@@ -7000,9 +7011,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
 [[package]]
 name = "tokenizers"
-version = "0.22.0"
+version = "0.22.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "af10f51be57162b69d90a15cb226eef12c9e4faecbd5e3ea98a86bfb920b3d71"
+checksum = "6475a27088c98ea96d00b39a9ddfb63780d1ad4cceb6f48374349a96ab2b7842"
 dependencies = [
  "ahash",
  "aho-corasick",
@@ -7454,9 +7465,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.101"
+version = "0.2.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e14915cadd45b529bb8d1f343c4ed0ac1de926144b746e2710f9cd05df6603b"
+checksum = "ab10a69fbd0a177f5f649ad4d8d3305499c42bab9aef2f7ff592d0ec8f833819"
 dependencies = [
  "cfg-if",
  "once_cell",
@@ -7467,9 +7478,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-backend"
-version = "0.2.101"
+version = "0.2.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e28d1ba982ca7923fd01448d5c30c6864d0a14109560296a162f80f305fb93bb"
+checksum = "0bb702423545a6007bbc368fde243ba47ca275e549c8a28617f56f6ba53b1d1c"
 dependencies = [
  "bumpalo",
  "log",
@@ -7481,9 +7492,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-futures"
-version = "0.4.51"
+version = "0.4.53"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0ca85039a9b469b38336411d6d6ced91f3fc87109a2a27b0c197663f5144dffe"
+checksum = "a0b221ff421256839509adbb55998214a70d829d3a28c69b4a6672e9d2a42f67"
 dependencies = [
  "cfg-if",
  "js-sys",
@@ -7494,9 +7505,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.101"
+version = "0.2.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7c3d463ae3eff775b0c45df9da45d68837702ac35af998361e2c84e7c5ec1b0d"
+checksum = "fc65f4f411d91494355917b605e1480033152658d71f722a90647f56a70c88a0"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -7504,9 +7515,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.101"
+version = "0.2.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7bb4ce89b08211f923caf51d527662b75bdc9c9c7aab40f86dcb9fb85ac552aa"
+checksum = "ffc003a991398a8ee604a401e194b6b3a39677b3173d6e74495eb51b82e99a32"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -7517,9 +7528,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.101"
+version = "0.2.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f143854a3b13752c6950862c906306adb27c7e839f7414cec8fea35beab624c1"
+checksum = "293c37f4efa430ca14db3721dfbe48d8c33308096bd44d80ebaa775ab71ba1cf"
 dependencies = [
  "unicode-ident",
 ]
@@ -7539,9 +7550,9 @@ dependencies = [
 
 [[package]]
 name = "web-sys"
-version = "0.3.78"
+version = "0.3.80"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77e4b637749ff0d92b8fad63aa1f7cff3cbe125fd49c175cd6345e7272638b12"
+checksum = "fbe734895e869dc429d78c4b433f8d17d95f8d05317440b4fad5ab2d33e596dc"
 dependencies = [
  "js-sys",
  "wasm-bindgen",

{spnl-0.6.0 → spnl-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spnl
-Version: 0.6.0
+Version: 0.6.1
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{spnl-0.6.0 → spnl-0.6.1}/spnl/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "spnl"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2024"
 license = "Apache-2.0"
 readme = "README.md"

{spnl-0.6.0 → spnl-0.6.1}/spnl/src/tokenize.rs

@@ -14,6 +14,7 @@ struct Tokenizer {
     cross_token: Option<u32>,
     plus_token: Option<u32>,
     block_size: usize,
+    assistant_suffix_num_tokens: usize,
 }
 
 impl Tokenizer {
@@ -33,8 +34,14 @@ impl Tokenizer {
     }
 
     fn assistanttok(&self, m: &str, tokens: &mut Vec<u32>) -> tokenizers::tokenizer::Result<()> {
+        let encoding = self.tok.encode_fast(self.assistant(m), false)?;
+        let extra = encoding.get_ids();
+
         self.extend_crop(
-            self.tok.encode_fast(self.assistant(m), false)?.get_ids(),
+            // TODO: for now, we always drop any assistant suffix. we
+            // will need to figure out how to isolatge these on their
+            // own block
+            &extra[0..extra.len() - self.assistant_suffix_num_tokens],
             tokens,
         );
         Ok(())
@@ -75,7 +82,7 @@ impl Tokenizer {
         // `tokens.len()+self.block_size-1`.
         let end = extra.len() + tokens.len();
         let nearest_block_boundary = end / self.block_size * self.block_size;
-        let amount_to_crop = end - nearest_block_boundary;
+        let amount_to_crop = ::std::cmp::min(extra.len(), end - nearest_block_boundary);
         let extra_end = extra.len() - amount_to_crop;
 
         self.extend(&extra[0..extra_end], tokens);
@@ -106,13 +113,44 @@ impl TokenizerState {
         block_size: usize,
     ) -> Result<::std::sync::Arc<Tokenizer>, ::std::sync::Arc<tokenizers::tokenizer::Error>> {
         self.cache.try_get_with(model.clone(), || {
+            let tok = tokenizers::tokenizer::Tokenizer::from_pretrained(model, None)?;
+            let tmpl = chat_template::detect(model)?;
+
+            let m = "hello";
+            let binding = tok.encode_fast(
+                chat_template::apply(tmpl, &[Message::Assistant(m.to_owned())], false),
+                false,
+            )?;
+            let binding2 = tok.encode_fast(m, false)?;
+            let with_chat_template = binding.get_ids();
+            let without_chat_template = binding2.get_ids();
+
+            // TODO this is imperfect...
+            let start_of_message_idx = with_chat_template
+                .iter()
+                .position(|t| *t == without_chat_template[0]);
+            let end_of_message_idx = start_of_message_idx
+                .map(|start_of_message_idx| start_of_message_idx + without_chat_template.len());
+            // [pppppmmmmmmmmmss]  <- ppppp are the prefix speical tokens added by chat template; ss suffix special tokens
+            //       ^ start_of_message_idx
+            //                ^ end_of_message_idx
+            let assistant_suffix_num_tokens = if let Some(end_of_message_idx) = end_of_message_idx {
+                with_chat_template.len() - end_of_message_idx
+            } else {
+                eprintln!(
+                    "Warning: could not determine length of end of assistant special token sequence"
+                );
+                0
+            };
+
             Ok(::std::sync::Arc::new(Tokenizer {
-                tmpl: chat_template::detect(model)?,
-                tok: tokenizers::tokenizer::Tokenizer::from_pretrained(model, None)?,
+                tmpl,
+                tok,
                 pad_token,
                 cross_token,
                 plus_token,
                 block_size,
+                assistant_suffix_num_tokens,
             }))
         })
     }
@@ -382,15 +420,9 @@ pub fn tokenize_query(
         temperature,
     } = query.g;
 
-    let s = ::std::time::Instant::now();
     let tok = state
         .get_or_create(&model, pad_token, cross_token, plus_token, block_size)
         .map_err(handle_arc_err)?;
-    println!(
-        "Spnl tokenize_query from pretrained {model}. Loaded in {:?}",
-        s.elapsed()
-    );
-
     let mut tokens: Vec<u32> = vec![];
     tokenize_part(&input, &tok, &mut tokens)
         .and_then(|()| add_final_assistant_token(&tok, &mut tokens))
@@ -460,3 +492,94 @@ pub fn tokenize_prepare(
         _ => todo!(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use itertools::Itertools;
+
+    const PAD_TOKEN: u32 = 27;
+    const BLOCK_SIZE: usize = 16;
+
+    const MODEL: &str = "ibm-granite/granite-3.3-2b-instruct"; // TODO find smaller model with public tokenizers.json
+    const START_OF_ROLE: u32 = 49152;
+    const END_OF_ROLE: u32 = 49153;
+    const END_OF_TEXT: u32 = 0;
+    const USER: u32 = 496;
+    const ASSISTANT: u32 = 17594;
+    const HELLO: u32 = 7656;
+    const LONGER: u32 = 8928;
+
+    fn tok() -> Result<::std::sync::Arc<Tokenizer>, ::std::sync::Arc<tokenizers::tokenizer::Error>>
+    {
+        init(2).get_or_create(&MODEL.into(), PAD_TOKEN, None, None, BLOCK_SIZE)
+    }
+
+    #[test]
+    fn create_tokenizer() -> Result<(), ::std::sync::Arc<tokenizers::tokenizer::Error>> {
+        tok().map(|_| ())
+    }
+
+    #[test]
+    fn user() -> Result<(), ::std::sync::Arc<tokenizers::tokenizer::Error>> {
+        assert_eq!(
+            tok().map(|tok| tok.user("hello"))?,
+            "<|start_of_role|>user<|end_of_role|>hello<|end_of_text|>"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn usertok() -> Result<(), ::std::sync::Arc<tokenizers::tokenizer::Error>> {
+        let mut tokens = vec![];
+        tok()?.usertok("hello", &mut tokens)?;
+        assert_eq!(
+            tokens,
+            [START_OF_ROLE, USER, END_OF_ROLE, HELLO, END_OF_TEXT]
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn assistant() -> Result<(), ::std::sync::Arc<tokenizers::tokenizer::Error>> {
+        assert_eq!(
+            tok().map(|tok| tok.assistant("hello"))?,
+            "<|start_of_role|>assistant<|end_of_role|>hello<|end_of_text|>"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn assistanttok_fully_cropped() -> Result<(), ::std::sync::Arc<tokenizers::tokenizer::Error>> {
+        let mut tokens = vec![];
+        tok()?.assistanttok("hello", &mut tokens)?;
+        let empty: &[u32] = &[];
+        assert_eq!(tokens, empty);
+        Ok(())
+    }
+
+    #[test]
+    fn assistanttok_partially_cropped() -> Result<(), ::std::sync::Arc<tokenizers::tokenizer::Error>>
+    {
+        let repeat_input = 17; // repeat this many times for the input message
+        let repeat_output = 12; // expect this many repetitions after cropping
+        let mut tokens = vec![];
+        tok()?.assistanttok(
+            format!(
+                "hello {}",
+                ::std::iter::repeat_n("longer", repeat_input).join(" ")
+            )
+            .as_str(),
+            &mut tokens,
+        )?;
+        assert_eq!(
+            tokens,
+            [START_OF_ROLE, ASSISTANT, END_OF_ROLE, HELLO]
+                .into_iter()
+                .chain(::std::iter::repeat_n(LONGER, repeat_output))
+                // .chain([END_OF_TEXT])
+                .collect::<Vec<u32>>(),
+        );
+        Ok(())
+    }
+}