embed precomputed odht tables for prefixes_of_mergeable_ranks (#13)

this provies a ~7% increase in count_tokens performance on intel
processors

the odht table lives in shared memory inside the mmap'd .so file, so
memory usage is improved as well

also speeds up runtime boot, since the table is now pregenerated

Author: tmm1

Cargo.lock

@@ -3,6 +3,7 @@ name = "tiktoken"
 version = "0.1.0"
 edition = "2021"
 rust-version = "1.57.0"
+build = "build.rs"
 
 [lib]
 name = "tiktoken"
@@ -13,6 +14,14 @@ tiktoken-rs = "0.5.4"
 memory-stats = "1.2.0"
 test-case = "2.0.0-rc3"
 
+[build-dependencies]
+hex = "0.4.3"
+sha2 = "0.10.6"
+base64 = "0.21.0"
+rustc-hash = "1.1.0"
+odht = "0.3.1"
+const-primes = "0.8.7"
+
 [dependencies]
 fancy-regex = "0.10.0"
 regex = "1.7.0"
@@ -23,6 +32,7 @@ sha2 = "0.10.6"
 base64 = "0.21.0"
 thiserror = "1.0.38"
 const-primes = "0.8.7"
+odht = "0.3.1"
 
 [[bench]]
 name = "bench"

Cargo.toml

@@ -0,0 +1,102 @@
+include!("src/load.rs");
+include!("src/rollhash.rs");
+include!("src/odht.rs");
+
+use std::env;
+use std::fs::File;
+use std::io::Write;
+use std::path::Path;
+use odht::HashTableOwned;
+use rustc_hash::FxHashSet as HashSet;
+
+fn main() {
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-changed=src/load.rs");
+
+    let out_dir = env::var("OUT_DIR").unwrap();
+    let mut file = File::create(&Path::new(&out_dir).join("static.rs")).unwrap();
+    writeln!(file, "pub mod data {{").unwrap();
+
+    generate("r50k_base",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/r50k_base.tiktoken"),
+            "306cd27f03c1a714eca7108e03d66b7dc042abe8c258b44c199a7ed9838dd930",
+        ).unwrap());
+
+    generate("p50k_base",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/p50k_base.tiktoken"),
+            "94b5ca7dff4d00767bc256fdd1b27e5b17361d7b8a5f968547f9f23eb70d2069",
+        ).unwrap());
+
+    generate("cl100k_base",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/cl100k_base.tiktoken"),
+            "223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7",
+        ).unwrap());
+
+    generate("o200k_base",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/o200k_base.tiktoken"),
+            "446a9538cb6c348e3516120d7c08b09f57c36495e2acfffe59a5bf8b0cfb1a2d",
+        ).unwrap());
+
+    generate("codestral",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/codestral.tiktoken"),
+            "bd5e66af07259851e88c3e483f88371dc2408cb0ce8b9787d29eaecdbb78eade",
+        ).unwrap());
+
+    generate("llama3",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/llama3.tiktoken"),
+            "82e9d31979e92ab929cd544440f129d9ecd797b69e327f80f17e1c50d5551b55",
+        ).unwrap());
+
+    generate("deepseekv2",
+        &mut file,
+        &load_tiktoken_bpe(
+            include_bytes!("data/deepseekv2.tiktoken"),
+            "3516b4e6e24389f7d1b288d861ce063da13296f916d29384e56ea9e0f6ba6674",
+        ).unwrap());
+
+    writeln!(file, "}}").unwrap();
+}
+
+fn generate(
+    name: &str,
+    file: &mut File,
+    mergeable_ranks: &HashMap<Vec<u8>, usize>,
+) {
+    writeln!(
+        file,
+        "    pub const {}_PREFIXES_ODHT: &'static [u8] = include_bytes!(\"{}.prefixes.odht\");",
+        name.to_uppercase(),
+        name
+    ).unwrap();
+
+    let mut prefixes_of_mergeable_ranks = mergeable_ranks
+    .keys()
+    .flat_map(|bytes| {
+        (1..=bytes.len())
+            .map(|i| roll_hash_slice(&bytes[..i]))
+            .collect::<Vec<_>>()
+    })
+    .collect::<HashSet<_>>();
+    prefixes_of_mergeable_ranks.insert(0);
+    prefixes_of_mergeable_ranks.shrink_to_fit();
+
+    let mut odht = HashTableOwned::<PrefixConfig>::with_capacity(prefixes_of_mergeable_ranks.len(), 50);
+    for prefix in prefixes_of_mergeable_ranks {
+        odht.insert(&prefix, &());
+    }
+
+    let mut file = File::create(format!("{}/{}.prefixes.odht", env::var("OUT_DIR").unwrap(), name)).unwrap();
+    file.write_all(odht.raw_bytes()).unwrap();
+}

build.rs

@@ -4,8 +4,12 @@ use rustc_hash::FxHashMap as HashMap;
 use rustc_hash::FxHashSet as HashSet;
 use std::sync::Arc;
 use thiserror::Error;
+use odht::HashTableOwned;
 use crate::rollhash::{roll_hash, roll_hash_slice};
 
+include!("odht.rs");
+include!(concat!(env!("OUT_DIR"), "/static.rs"));
+
 /// A struct that represents an encoding scheme based on byte-pair encoding (BPE).
 #[derive(Debug)]
 pub struct Encoding {
@@ -16,7 +20,7 @@ pub struct Encoding {
     /// The maximum length of the keys in `mergeable_ranks`.
     mergeable_ranks_max_key_len: usize,
     /// All prefixes of the mergeable ranks. May or may not be tokens themselves!
-    prefixes_of_mergeable_ranks: HashSet<i64>,
+    prefixes_of_mergeable_ranks: HashTableOwned<PrefixConfig>,
     /// The map from special token strings to their values.
     special_tokens: HashMap<String, usize>,
     /// The maximum token value in the encoding.
@@ -97,16 +101,18 @@ impl Encoding {
         )
         .map_err(|e| EncodingError::GenericEncodingError(format!("Error creating core BPE: {}", e)))?;
 
-        let mut prefixes_of_mergeable_ranks = mergeable_ranks
-            .keys()
-            .flat_map(|bytes| {
-                (1..=bytes.len())
-                    .map(|i| roll_hash_slice(&bytes[..i]))
-                    .collect::<Vec<_>>()
+        let prefixes_of_mergeable_ranks = unsafe {
+            HashTableOwned::<PrefixConfig>::from_raw_bytes_unchecked(match name {
+                "r50k_base" => data::R50K_BASE_PREFIXES_ODHT,
+                "p50k_base" => data::P50K_BASE_PREFIXES_ODHT,
+                "cl100k_base" => data::CL100K_BASE_PREFIXES_ODHT,
+                "o200k_base" => data::O200K_BASE_PREFIXES_ODHT,
+                "codestral" => data::CODESTRAL_PREFIXES_ODHT,
+                "llama3" => data::LLAMA3_PREFIXES_ODHT,
+                "deepseekv2" => data::DEEPSEEKV2_PREFIXES_ODHT,
+                _ => return Err(EncodingError::GenericEncodingError(format!("Embedded prefix table not found for encoding: {}", name))),
             })
-            .collect::<HashSet<_>>();
-        prefixes_of_mergeable_ranks.insert(0);
-        prefixes_of_mergeable_ranks.shrink_to_fit();
+        };
 
         Ok(Self {
             name: name.to_string(),
@@ -148,7 +154,7 @@ impl Encoding {
             // or if the current token is not in the prefixes of mergeable ranks,
             // we need to split the current token and begin actually checking for the largest
             // mergeable prefix
-            while !self.prefixes_of_mergeable_ranks.contains(&current_token_hash)
+            while !self.prefixes_of_mergeable_ranks.contains_key(&current_token_hash)
                 || current_token.len() > self.mergeable_ranks_max_key_len
             {
                 if current_token.len() > 1 {

src/encoding.rs

@@ -0,0 +1,20 @@
+use odht::{Config, FxHashFn};
+
+pub struct PrefixConfig;
+
+impl Config for PrefixConfig {
+    type Key = i64;
+    type Value = ();
+    type EncodedKey = [u8; 8];
+    type EncodedValue = [u8; 0];
+    type H = FxHashFn;
+
+    #[inline(always)]
+    fn encode_key(k: &Self::Key) -> Self::EncodedKey { k.to_le_bytes() }
+    #[inline(always)]
+    fn encode_value(_: &Self::Value) -> Self::EncodedValue { [] }
+    #[inline(always)]
+    fn decode_key(k: &Self::EncodedKey) -> Self::Key { i64::from_le_bytes(*k) }
+    #[inline(always)]
+    fn decode_value(_: &Self::EncodedValue) -> Self::Value { () }
+}