Backtrack tmp.

Author: Narsil

bindings/python/benches/test_backtrack.py

@@ -0,0 +1,88 @@
+import os
+import argparse
+import datetime
+from datasets import load_dataset
+from tokenizers import Tokenizer
+from typing import Tuple
+
+MODEL_ID = "meta-llama/Meta-Llama-3.1-8B"
+DATASET = "facebook/xnli"
+DATASET_CONFIG = "all_languages"
+DEFAULT_THREADS = [2**i for i in range(8) if 2**i <= os.cpu_count()]
+
+
+def format_byte_size(num_bytes: int) -> Tuple[str, str]:
+    """Convert bytes to a human-readable format (KB, MB, GB)."""
+    num_bytes_f = float(num_bytes)
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if num_bytes_f < 1024:
+            return f"{num_bytes_f:.2f} {unit}", unit
+        num_bytes_f /= 1024
+    return f"{num_bytes_f:.2f} PB", "PB"
+
+
+def test(model: str, dataset: str, dataset_config: str):
+    dataset_xnli = load_dataset(dataset, dataset_config)
+    tokenizer = Tokenizer.from_pretrained(model)
+    tokenizer2 = Tokenizer.from_pretrained(model)
+    tokenizer2.enable_backtrack()
+
+    for easy in ["1880", " cream"]:
+        encoded = tokenizer.encode(easy)
+        encoded2 = tokenizer2.encode(easy)
+        if encoded.ids != encoded2.ids:
+            import ipdb
+
+            ipdb.set_trace()
+        assert encoded.ids == encoded2.ids
+
+    sentences = []
+    en_sentences = []
+    for _i, item in enumerate(dataset_xnli["train"]):
+        # sentence = item["premise"]["en"]
+        # sentences.append(sentence)
+        for lang, sentence in item["premise"].items():
+            if lang == "en":
+                en_sentences.append(sentence)
+            sentences.append(sentence)
+    sentences = en_sentences + sentences
+
+    start = datetime.datetime.now()
+    encoded = tokenizer.encode_batch_fast(sentences)
+    print(f"Took {datetime.datetime.now() - start}")
+
+    start = datetime.datetime.now()
+    encoded2 = tokenizer2.encode_batch_fast(sentences)
+    print(f"Took {datetime.datetime.now() - start}")
+
+    assert len(encoded) == len(encoded2)
+    assert len(encoded) == len(sentences)
+    total = 0
+    correct = 0
+    for enc, enc2, sentence in zip(encoded, encoded2, sentences):
+        # if enc.ids != enc2.ids:
+        #     print(enc.ids)
+        #     print(enc2.ids)
+        if enc.ids == enc2.ids:
+            correct += 1
+        total += 1
+        assert enc.ids == enc2.ids, f"{enc.ids} != {enc2.ids} (Source: {sentence}"
+    print(f"{correct} / {total} ({correct / total * 100:.2f}%%)")
+    # print("All good !")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="bench_tokenizer",
+        description="Getting a feel for speed when tokenizing",
+    )
+    parser.add_argument("-m", "--model", default=MODEL_ID, type=str)
+    parser.add_argument("-d", "--dataset", default=DATASET, type=str)
+    parser.add_argument("-ds", "--dataset-config", default=DATASET_CONFIG, type=str)
+    args = parser.parse_args()
+    test(args.model, args.dataset, args.dataset_config)
+
+
+# Call the function to run the benchmark
+if __name__ == "__main__":
+    main()

bindings/python/src/tokenizer.rs

@@ -1,6 +1,8 @@
 use serde::Serialize;
 use std::collections::{hash_map::DefaultHasher, HashMap};
 use std::hash::{Hash, Hasher};
+use tk::pre_tokenizers::byte_level::ByteLevel;
+use tk::ModelWrapper;
 
 use numpy::{npyffi, PyArray1, PyArrayMethods};
 use pyo3::class::basic::CompareOp;
@@ -1118,6 +1120,19 @@ impl PyTokenizer {
             .into()
         })
     }
+    ///
+    #[pyo3(signature = ())]
+    #[pyo3(text_signature = "(self)")]
+    fn enable_backtrack(&mut self) -> PyResult<()> {
+        // self.tokenizer.with_pre_tokenizer(None::<ByteLevel>);
+        let model = self.tokenizer.get_model();
+        let mut model = model.model.write().unwrap();
+        let ModelWrapper::BPE(ref mut model) = *model else {
+            todo!();
+        };
+        model.enable_backtrack();
+        Ok(())
+    }
 
     /// Decode the given list of ids back to a string
     ///

tokenizers/Cargo.toml

@@ -69,6 +69,8 @@ monostate = "0.1.12"
 ahash = { version = "0.8.11", features = ["serde"] }
 dary_heap = { version = "0.3.6", features = ["serde"] }
 compact_str = { version = "0.9", features = ["serde"] }
+fnv = "1.0.7"
+aneubeck-daachorse = "1.1.1"
 
 [features]
 default = ["progressbar", "onig", "esaxx_fast"]