Consolidated optimization ahash dary compact str (#1799)

* free speed/mem optimizations with ahash, dary_heap, and compact_str (bindings broken without library refactor)

* Rebased ahash.

* Removing data files.

* Fixing (dummily) the python pyobject conversion.

* Fixing the surface by not providing ahash.

* Fixing the python side with the ahash public api removal.

* Cleanup.

* Removing test file.

* Remove dead file.

* Bad conflict resolution.

* Bad merge 2.

Author: Narsil

bindings/node/Cargo.toml

@@ -14,6 +14,7 @@ napi        = "2"
 napi-derive = "2"
 serde       = { version = "1.0.163", features = ["derive"] }
 tokenizers  = { path = "../../tokenizers/" }
+ahash = { version = "0.8.11", features = ["serde"] }
 
 [build-dependencies]
 napi-build = "2"

bindings/node/src/models.rs

@@ -1,14 +1,15 @@
 use crate::arc_rwlock_serde;
 use crate::tasks::models::{BPEFromFilesTask, WordLevelFromFilesTask, WordPieceFromFilesTask};
 use crate::trainers::Trainer;
+use ahash::AHashMap;
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, RwLock};
 use tokenizers as tk;
-use tokenizers::models::bpe::{BpeBuilder, Merges, Vocab};
+use tokenizers::models::bpe::{BpeBuilder, Merges};
 use tokenizers::models::wordlevel::WordLevelBuilder;
 use tokenizers::models::wordpiece::WordPieceBuilder;
 
@@ -44,8 +45,13 @@ impl Bpe {
   }
 
   #[napi(factory, ts_return_type = "Model")]
-  pub fn init(vocab: Vocab, merges: Merges, options: Option<BpeOptions>) -> Result<Model> {
+  pub fn init(
+    vocab: HashMap<String, u32>,
+    merges: Merges,
+    options: Option<BpeOptions>,
+  ) -> Result<Model> {
     let options = options.unwrap_or_default();
+    let vocab: AHashMap<_, _> = vocab.into_iter().collect();
     let mut builder = tk::models::bpe::BPE::builder().vocab_and_merges(vocab, merges);
     builder = options.apply_to_bpe_builder(builder);
     let model = builder
@@ -206,10 +212,11 @@ pub struct WordPiece {}
 #[napi]
 impl WordPiece {
   #[napi(factory, ts_return_type = "Model")]
-  pub fn init(vocab: Vocab, options: Option<WordPieceOptions>) -> Result<Model> {
+  pub fn init(vocab: HashMap<String, u32>, options: Option<WordPieceOptions>) -> Result<Model> {
     let options = options.unwrap_or_default();
 
-    let mut builder = tk::models::wordpiece::WordPiece::builder().vocab(vocab);
+    let mut builder = tk::models::wordpiece::WordPiece::builder()
+      .vocab(vocab.into_iter().collect::<AHashMap<_, _>>());
     builder = options.apply_to_wordpiece_builder(builder);
     let model = builder
       .build()
@@ -263,9 +270,10 @@ pub struct WordLevel {}
 #[napi]
 impl WordLevel {
   #[napi(factory, ts_return_type = "Model")]
-  pub fn init(vocab: Vocab, options: Option<WordLevelOptions>) -> Result<Model> {
+  pub fn init(vocab: HashMap<String, u32>, options: Option<WordLevelOptions>) -> Result<Model> {
     let options = options.unwrap_or_default();
-    let mut builder = tk::models::wordlevel::WordLevel::builder().vocab(vocab);
+    let mut builder =
+      tk::models::wordlevel::WordLevel::builder().vocab(vocab.into_iter().collect());
     builder = options.apply_to_wordlevel_builder(builder);
     let model = builder
       .build()

bindings/python/Cargo.toml

@@ -18,6 +18,7 @@ pyo3 = { version = "0.25", features = ["abi3", "abi3-py39", "py-clone"] }
 numpy = "0.25"
 ndarray = "0.16"
 itertools = "0.14"
+ahash = { version = "0.8.11", features = ["serde"] }
 
 [dependencies.tokenizers]
 path = "../../tokenizers"

bindings/python/src/models.rs

@@ -4,11 +4,12 @@ use std::sync::{Arc, RwLock};
 
 use crate::token::PyToken;
 use crate::trainers::PyTrainer;
+use ahash::AHashMap;
 use pyo3::exceptions;
 use pyo3::prelude::*;
 use pyo3::types::*;
 use serde::{Deserialize, Serialize};
-use tk::models::bpe::{BpeBuilder, Merges, Vocab, BPE};
+use tk::models::bpe::{BpeBuilder, Merges, BPE};
 use tk::models::unigram::Unigram;
 use tk::models::wordlevel::WordLevel;
 use tk::models::wordpiece::{WordPiece, WordPieceBuilder};
@@ -347,9 +348,10 @@ macro_rules! setter {
 
 #[derive(FromPyObject)]
 enum PyVocab {
-    Vocab(Vocab),
+    Vocab(HashMap<String, u32>),
     Filename(String),
 }
+
 #[derive(FromPyObject)]
 enum PyMerges {
     Merges(Merges),
@@ -454,6 +456,7 @@ impl PyBPE {
         if let (Some(vocab), Some(merges)) = (vocab, merges) {
             match (vocab, merges) {
                 (PyVocab::Vocab(vocab), PyMerges::Merges(merges)) => {
+                    let vocab: AHashMap<_, _> = vocab.into_iter().collect();
                     builder = builder.vocab_and_merges(vocab, merges);
                 }
                 (PyVocab::Filename(vocab_filename), PyMerges::Filename(merges_filename)) => {
@@ -494,13 +497,15 @@ impl PyBPE {
     ///         The vocabulary and merges loaded into memory
     #[staticmethod]
     #[pyo3(text_signature = "(self, vocab, merges)")]
-    fn read_file(vocab: &str, merges: &str) -> PyResult<(Vocab, Merges)> {
-        BPE::read_file(vocab, merges).map_err(|e| {
+    fn read_file(vocab: &str, merges: &str) -> PyResult<(HashMap<String, u32>, Merges)> {
+        let (vocab, merges) = BPE::read_file(vocab, merges).map_err(|e| {
             exceptions::PyException::new_err(format!(
                 "Error while reading vocab & merges files: {}",
                 e
             ))
-        })
+        })?;
+        let vocab = vocab.into_iter().collect();
+        Ok((vocab, merges))
     }
 
     /// Instantiate a BPE model from the given files.
@@ -536,6 +541,7 @@ impl PyBPE {
         let (vocab, merges) = BPE::read_file(vocab, merges).map_err(|e| {
             exceptions::PyException::new_err(format!("Error while reading BPE files: {}", e))
         })?;
+        let vocab = vocab.into_iter().collect();
         Py::new(
             py,
             PyBPE::new(
@@ -668,6 +674,7 @@ impl PyWordPiece {
         if let Some(vocab) = vocab {
             match vocab {
                 PyVocab::Vocab(vocab) => {
+                    let vocab: AHashMap<_, _> = vocab.into_iter().collect();
                     builder = builder.vocab(vocab);
                 }
                 PyVocab::Filename(vocab_filename) => {
@@ -699,10 +706,11 @@ impl PyWordPiece {
     ///     :obj:`Dict[str, int]`: The vocabulary as a :obj:`dict`
     #[staticmethod]
     #[pyo3(text_signature = "(vocab)")]
-    fn read_file(vocab: &str) -> PyResult<Vocab> {
-        WordPiece::read_file(vocab).map_err(|e| {
+    fn read_file(vocab: &str) -> PyResult<HashMap<String, u32>> {
+        let vocab = WordPiece::read_file(vocab).map_err(|e| {
             exceptions::PyException::new_err(format!("Error while reading WordPiece file: {}", e))
-        })
+        })?;
+        Ok(vocab.into_iter().collect())
     }
 
     /// Instantiate a WordPiece model from the given file
@@ -734,6 +742,7 @@ impl PyWordPiece {
         let vocab = WordPiece::read_file(vocab).map_err(|e| {
             exceptions::PyException::new_err(format!("Error while reading WordPiece file: {}", e))
         })?;
+        let vocab = vocab.into_iter().collect();
         Py::new(
             py,
             PyWordPiece::new(py, Some(PyVocab::Vocab(vocab)), kwargs)?,
@@ -778,6 +787,7 @@ impl PyWordLevel {
         if let Some(vocab) = vocab {
             match vocab {
                 PyVocab::Vocab(vocab) => {
+                    let vocab = vocab.into_iter().collect();
                     builder = builder.vocab(vocab);
                 }
                 PyVocab::Filename(vocab_filename) => {
@@ -818,10 +828,12 @@ impl PyWordLevel {
     ///     :obj:`Dict[str, int]`: The vocabulary as a :obj:`dict`
     #[staticmethod]
     #[pyo3(text_signature = "(vocab)")]
-    fn read_file(vocab: &str) -> PyResult<Vocab> {
-        WordLevel::read_file(vocab).map_err(|e| {
+    fn read_file(vocab: &str) -> PyResult<HashMap<String, u32>> {
+        let vocab = WordLevel::read_file(vocab).map_err(|e| {
             exceptions::PyException::new_err(format!("Error while reading WordLevel file: {}", e))
-        })
+        })?;
+        let vocab: HashMap<_, _> = vocab.into_iter().collect();
+        Ok(vocab)
     }
 
     /// Instantiate a WordLevel model from the given file
@@ -853,6 +865,7 @@ impl PyWordLevel {
         let vocab = WordLevel::read_file(vocab).map_err(|e| {
             exceptions::PyException::new_err(format!("Error while reading WordLevel file: {}", e))
         })?;
+        let vocab = vocab.into_iter().collect();
         Py::new(
             py,
             PyWordLevel::new(py, Some(PyVocab::Vocab(vocab)), unk_token)?,

tokenizers/Cargo.toml

@@ -66,6 +66,9 @@ fancy-regex = { version = "0.14", optional = true}
 getrandom = { version = "0.3" }
 esaxx-rs = { version = "0.1.10", default-features = false, features=[]}
 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"] }
 
 [features]
 default = ["progressbar", "onig", "esaxx_fast"]

tokenizers/src/models/bpe/model.rs

@@ -2,19 +2,21 @@ use super::{super::OrderedVocabIter, trainer::BpeTrainer, Error, Pair, Word};
 use crate::tokenizer::{Model, Result, Token};
 use crate::utils::cache::{Cache, DEFAULT_CACHE_CAPACITY, MAX_LENGTH};
 use crate::utils::iter::ResultShunt;
+use ahash::AHashMap;
 use serde_json::Value;
 use std::borrow::Cow;
+
+use std::collections::HashMap;
 use std::{
-    collections::HashMap,
     fs::File,
     io::prelude::*,
     io::{BufRead, BufReader},
     path::{Path, PathBuf},
 };
 
-pub type Vocab = HashMap<String, u32>;
-type VocabR = HashMap<u32, String>;
-pub type MergeMap = HashMap<Pair, (u32, u32)>;
+pub type Vocab = AHashMap<String, u32>;
+type VocabR = AHashMap<u32, String>;
+pub type MergeMap = AHashMap<Pair, (u32, u32)>;
 pub type Merges = Vec<(String, String)>;
 
 struct Config {
@@ -41,7 +43,7 @@ impl Default for BpeBuilder {
         Self {
             config: Config {
                 files: None,
-                vocab: HashMap::new(),
+                vocab: AHashMap::new(),
                 merges: vec![],
                 cache_capacity: DEFAULT_CACHE_CAPACITY,
                 dropout: None,
@@ -71,8 +73,12 @@ impl BpeBuilder {
 
     /// Set the vocab (token -> ID) and merges mappings.
     #[must_use]
-    pub fn vocab_and_merges(mut self, vocab: Vocab, merges: Merges) -> Self {
-        self.config.vocab = vocab;
+    pub fn vocab_and_merges<V: Into<AHashMap<String, u32>>>(
+        mut self,
+        vocab: V,
+        merges: Merges,
+    ) -> Self {
+        self.config.vocab = vocab.into();
         self.config.merges = merges;
         self
     }
@@ -324,7 +330,7 @@ impl BPE {
         let mut buffer = String::new();
         vocab_file.read_to_string(&mut buffer)?;
         let json: Value = serde_json::from_str(&buffer)?;
-        let mut vocab = HashMap::new();
+        let mut vocab = AHashMap::new();
         match json {
             Value::Object(m) => {
                 for (token, id) in m {
@@ -361,8 +367,8 @@ impl BPE {
         }
     }
 
-    pub fn get_vocab(&self) -> Vocab {
-        self.vocab.clone()
+    pub fn get_vocab(&self) -> HashMap<String, u32> {
+        self.vocab.clone().into_iter().collect()
     }
 
     pub fn get_unk_token(&self) -> &Option<String> {
@@ -494,7 +500,7 @@ impl Model for BPE {
     type Trainer = BpeTrainer;
 
     fn get_vocab(&self) -> HashMap<String, u32> {
-        self.vocab.clone()
+        self.vocab.clone().into_iter().collect()
     }
 
     fn get_vocab_size(&self) -> usize {

tokenizers/src/models/bpe/serialization.rs

@@ -1,10 +1,10 @@
 use super::{super::OrderedVocabIter, convert_merges_to_hashmap, BpeBuilder, Pair, BPE};
+use ahash::AHashMap;
 use serde::{
     de::{Error, MapAccess, Visitor},
     ser::SerializeStruct,
     Deserialize, Deserializer, Serialize, Serializer,
 };
-use std::collections::HashMap;
 
 impl Serialize for BPE {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
@@ -80,7 +80,7 @@ impl<'de> Visitor<'de> for BPEVisitor {
         V: MapAccess<'de>,
     {
         let mut builder = BpeBuilder::new();
-        let mut vocab: Option<HashMap<String, u32>> = None;
+        let mut vocab: Option<AHashMap<String, u32>> = None;
 
         #[derive(Debug, Deserialize)]
         #[serde(untagged)]