Repository: huggingface/tokenizers (language: rust)

Issue:
"the" token is splitted to "t" "h" "e" in large scale corpus
I train the bpe (BERTWordPiece) by my own, using the RoBERTa 160GB data. However, I found all "the" is broken. 
I check the learned vocab.txt, and found "the" is not in there too.

The parameters used in training bpe.
```
bpe.train(args.inputs, vocab_size=32768, min_frequency=20)
```

I guess the count of `the` may overflow in the large-scale corpus.

an example

```
john . who . like me . is a writer with an interest in t ##h ##e interp ##lay of words and images . assumed i ' d hidden a cryptic message in t ##h ##e random sequence of little pictures . " write a package . spider earth cutting ? "
he asked . i only responded with five more random emoji . " t ##h ##e cheese running bear got caught in a rains ##how ##er looking for diamonds . " he muse ##d .
and thus began a game that has continued to this day .
one of us texts five random emoji to t ##h ##e other . and t ##h ##e recipient has to write a story ( roughly of tweet length ) to explain each of t ##h ##e emoji in order .
below are eight of our collaborative creations . 
james hannah ##am ' s most recent novel . delicious foods . won t ##h ##e pen / faulk ##ner and hurst ##on / wright legacy awards . and was a new york times notable book .
he teaches in t ##h ##e writing program at t ##h ##e pratt institute
```


Discussion / hints:
It's possible, I made a temporary branch with u64 everywhere instead of u32:

https://github.com/huggingface/tokenizers/tree/u64_branch

It's a temporary fix, we probably need to think a bit more about how to go about it in the general case.

You can install from source by going into `bindings/python`, and then `pip install -e .`
Thanks for the quick response. Actually, I made `the` as a special token as a workaround, since other high-frequency (like `and`) tokens seems ok.
If "th" and "he" are not tokens either you're probably going to have a subpar tokenizer
@Narsil Thanks, I just check it, and seems both `th` and `he` is there. 
I think it is better to fix the overflow problem, in case we have the larger corpus.
I didn't mean to close.
seems `u64` branch cannot fix the problem.
the error
```
thread '<unnamed>' panicked at 'Trainer should know how to build BPE: MergeTokenOutOfVocabulary("##h##e")', /home/xxx/tokenizers/tokenizers/src/models/bpe/trainer.rs:573:13
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Traceback (most recent call last):
  File "train_wp.py", line 55, in <module>
    main()
  File "train_wp.py", line 50, in main
    bpe.train(args.inputs, vocab_size=args.vocab_size, special_tokens=special_tokens,
  File "/home/xxx/tokenizers/bindings/python/py_src/tokenizers/implementations/bert_wordpiece.py", line 117, in train
    self._tokenizer.train(trainer, files)
pyo3_runtime.PanicException: Trainer should know how to build BPE: MergeTokenOutOfVocabulary("##h##e")
```
It seems the master branch also failed with the same error.
the same code in 0.8.1 could be run.

```
    bpe = BertWordPieceTokenizer(clean_text=True, strip_accents=True, lowercase=True)
    bpe.train(args.inputs)
```

updated:
it seems this commit (https://github.com/huggingface/tokenizers/commit/95cc8c47ad2fc1af13be406e3bdcdbd2f09fe71e) causes the problem.

master branch seems is much slower in file reading.

0.8.1
```
[00:03:13] Reading files (19521 Mo)                 ███████████████████████████████████████████████████████████████████████████████████████████████████████████                 100
[00:00:17] Tokenize words                           ███████████████████████████████████████████████████████████████████████████████████████████████████████████ 6593316  /  6593316
[00:00:14] Count pairs                              ███████████████████████████████████████████████████████████████████████████████████████████████████████████ 6593316  /  6593316
[00:00:58] Compute merges                           ███████████████████████████████████████████████████████████████████████████████████████████████████████████ 31116    /    31116

```

master
```
[00:17:46] Reading files (19521 Mo)                 ███████████████████████████████████████████████████████████████████████████████████████████████████████████                 100
[00:00:18] Tokenize words                           ███████████████████████████████████████████████████████████████████████████████████████████████████████████ 6593316  /  6593316
[00:00:15] Count pairs                              ███████████████████████████████████████████████████████████████████████████████████████████████████████████ 6593316  /  6593316
[00:01:13] Compute merges                           ███████████████████████████████████████████████████████████████████████████████████████████████████████████ 31117    /    31117

```
So there is definitely a slowdown with master that is supposed to be expected, that's linked to better (and correct) offsets tracking. However it should *not* be that bad, I'm going to try and replicate this. (It does not seem to be the case on smaller files ~500Mo, the slowdown is only ~20%)
This is the sort of variance I'm getting:

master:
```
[00:01:37] Reading files (1048 Mo)                  █████████████████████████████████████████████████████████████████████                 100
[00:00:02] Tokenize words                           █████████████████████████████████████████████████████████████████████ 1882027  /  1882027
[00:00:09] Count pairs                              █████████████████████████████████████████████████████████████████████ 1882027  /  1882027
[00:00:10] Compute merges                           █████████████████████████████████████████████████████████████████████ 28375    /    28375
```

v.0.8.1
```
[00:01:20] Reading files (1048 Mo)                  █████████████████████████████████████████████████████████████████████                 100
[00:00:03] Tokenize words                           █████████████████████████████████████████████████████████████████████ 1882027  /  1882027
[00:00:10] Count pairs                              █████████████████████████████████████████████████████████████████████ 1882027  /  1882027
[00:00:11] Compute merges                           █████████████████████████████████████████████████████████████████████ 28374    /    28374
```

Are you sure your master is up to date ? This kind of slowdown you're experiencing looks like *debug* build kind of problem.

Code used:
```python
from tokenizers import BertWordPieceTokenizer

bpe = BertWordPieceTokenizer(clean_text=True, strip_accents=True, lowercase=False)
bpe.train(["out1.txt"])

print(bpe.encode("I love this [SEP]").ids)
bpe.save("tokenizer.json", pretty=True)
```
I used this commit https://github.com/huggingface/tokenizers/commit/36832bfa1292b298c36447d7d2f5a3366c26644c, as the latter of it cannot run. 
I also set RAYON_RS_NUM_CPUS=16 , maybe the multi-threading performance is down?
Has there been any progress on integrating this into main? We are running into the same issue with the count of many two-character tokens overflowing and not ending up in the final vocab with a sufficiently large text corpus.
@gpauloski Not yet, the tentative branch  proposed here is super old.

You're welcome to attempt a PR, but making everything templated for both  `u32` and `u64` is not that trivial in Rust (and we need the logic to use `u64` only when necessary.

So far the recommendation has been to use something like `spm` which can handle `u64` for training, and converting it to `tokenizers` (this is what happened for `bigscience`)
Thanks, @Narsil! I'll try `spm`.
Actually keeping the tokenizer training data under `u32` is the final route that was taken for `bigscience` (but the spm u64 was also used when trying to figure out the best tokenizer)

Fix patch:
diff --git a/bindings/python/src/trainers.rs b/bindings/python/src/trainers.rs
index 1c1c9310a..707dc7230 100644
--- a/bindings/python/src/trainers.rs
+++ b/bindings/python/src/trainers.rs
@@ -183,12 +183,12 @@ impl PyBpeTrainer {
     }
 
     #[getter]
-    fn get_min_frequency(self_: PyRef<Self>) -> u32 {
+    fn get_min_frequency(self_: PyRef<Self>) -> u64 {
         getter!(self_, BpeTrainer, min_frequency)
     }
 
     #[setter]
-    fn set_min_frequency(self_: PyRef<Self>, freq: u32) {
+    fn set_min_frequency(self_: PyRef<Self>, freq: u64) {
         setter!(self_, BpeTrainer, min_frequency, freq);
     }
 
@@ -397,12 +397,12 @@ impl PyWordPieceTrainer {
     }
 
     #[getter]
-    fn get_min_frequency(self_: PyRef<Self>) -> u32 {
+    fn get_min_frequency(self_: PyRef<Self>) -> u64 {
         getter!(self_, WordPieceTrainer, min_frequency())
     }
 
     #[setter]
-    fn set_min_frequency(self_: PyRef<Self>, freq: u32) {
+    fn set_min_frequency(self_: PyRef<Self>, freq: u64) {
         setter!(self_, WordPieceTrainer, @set_min_frequency, freq);
     }
 
@@ -589,12 +589,12 @@ impl PyWordLevelTrainer {
     }
 
     #[getter]
-    fn get_min_frequency(self_: PyRef<Self>) -> u32 {
+    fn get_min_frequency(self_: PyRef<Self>) -> u64 {
         getter!(self_, WordLevelTrainer, min_frequency)
     }
 
     #[setter]
-    fn set_min_frequency(self_: PyRef<Self>, freq: u32) {
+    fn set_min_frequency(self_: PyRef<Self>, freq: u64) {
         setter!(self_, WordLevelTrainer, min_frequency, freq);
     }
 
diff --git a/tokenizers/src/models/bpe/trainer.rs b/tokenizers/src/models/bpe/trainer.rs
index 3821cdab4..303fdbc81 100644
--- a/tokenizers/src/models/bpe/trainer.rs
+++ b/tokenizers/src/models/bpe/trainer.rs
@@ -11,7 +11,7 @@ use std::collections::{BinaryHeap, HashMap, HashSet};
 #[derive(Debug, Eq)]
 struct Merge {
     pair: Pair,
-    count: u32,
+    count: u64,
     pos: HashSet<usize>,
 }
 impl PartialEq for Merge {
@@ -36,7 +36,7 @@ impl Ord for Merge {
 }
 
 struct Config {
-    min_frequency: u32,
+    min_frequency: u64,
     vocab_size: usize,
     show_progress: bool,
     special_tokens: Vec<AddedToken>,
@@ -79,7 +79,7 @@ impl BpeTrainerBuilder {
 
     /// Set the expected minimum frequency
     #[must_use]
-    pub fn min_frequency(mut self, frequency: u32) -> Self {
+    pub fn min_frequency(mut self, frequency: u64) -> Self {
         self.config.min_frequency = frequency;
         self
     }
@@ -176,7 +176,7 @@ impl BpeTrainerBuilder {
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq)]
 pub struct BpeTrainer {
     /// The minimum frequency a pair must have to produce a merge operation
-    pub min_frequency: u32,
+    pub min_frequency: u64,
     /// The target vocabulary size
     pub vocab_size: usize,
     /// Whether to show progress while training
@@ -195,7 +195,7 @@ pub struct BpeTrainer {
     /// An optional parameter to limit the max length of any single token
     pub max_token_length: Option<usize>,
 
-    words: HashMap<String, u32>,
+    words: HashMap<String, u64>,
 }
 
 impl Default for BpeTrainer {
@@ -205,7 +205,7 @@ impl Default for BpeTrainer {
 }
 
 impl BpeTrainer {
-    pub fn new(min_frequency: u32, vocab_size: usize) -> Self {
+    pub fn new(min_frequency: u64, vocab_size: usize) -> Self {
         Self {
             min_frequency,
             vocab_size,
@@ -263,7 +263,7 @@ impl BpeTrainer {
     /// Compute the initial alphabet and limit it if relevant
     fn compute_alphabet(
         &self,
-        wc: &HashMap<String, u32>,
+        wc: &HashMap<String, u64>,
         w2id: &mut HashMap<String, u32>,
         id2w: &mut Vec<String>,
     ) {
@@ -322,13 +322,13 @@ impl BpeTrainer {
     /// Tokenize words and add subwords to the vocabulary when relevant
     fn tokenize_words(
         &self,
-        wc: &HashMap<String, u32>,
+        wc: &HashMap<String, u64>,
         w2id: &mut HashMap<String, u32>,
         id2w: &mut Vec<String>,
         p: &Option<ProgressBar>,
-    ) -> (Vec<Word>, Vec<u32>) {
+    ) -> (Vec<Word>, Vec<u64>) {
         let mut words: Vec<Word> = Vec::with_capacity(wc.len());
-        let mut counts: Vec<u32> = Vec::with_capacity(wc.len());
+        let mut counts: Vec<u64> = Vec::with_capacity(wc.len());
 
         for (word, count) in wc {
             let mut current_word = Word::new();
@@ -373,7 +373,7 @@ impl BpeTrainer {
     fn count_pairs(
         &self,
         words: &[Word],
-        counts: &[u32],
+        counts: &[u64],
         p: &Option<ProgressBar>,
     ) -> (HashMap<Pair, i32>, HashMap<Pair, HashSet<usize>>) {
         words
@@ -431,7 +431,7 @@ impl BpeTrainer {
 
     pub fn do_train(
         &self,
-        word_counts: &HashMap<String, u32>,
+        word_counts: &HashMap<String, u64>,
         model: &mut BPE,
     ) -> Result<Vec<AddedToken>> {
         let mut word_to_id: HashMap<String, u32> = HashMap::with_capacity(self.vocab_size);
@@ -470,7 +470,7 @@ impl BpeTrainer {
             if count > 0 {
                 queue.push(Merge {
                     pair,
-                    count: count as u32,
+                    count: count as u64,
                     pos,
                 });
             }
@@ -493,8 +493,8 @@ impl BpeTrainer {
             }
 
             let mut top = queue.pop().unwrap();
-            if top.count != pair_counts[&top.pair] as u32 {
-                top.count = pair_counts[&top.pair] as u32;
+            if top.count != pair_counts[&top.pair] as u64 {
+                top.count = pair_counts[&top.pair] as u64;
                 queue.push(top);
                 continue;
             }
@@ -573,7 +573,7 @@ impl BpeTrainer {
                 if count > 0 {
                     queue.push(Merge {
                         pair,
-                        count: count as u32,
+                        count: count as u64,
                         pos,
                     });
                 }
@@ -632,7 +632,7 @@ impl Trainer for BpeTrainer {
         S: AsRef<str> + Send,
         F: Fn(&str) -> Result<Vec<String>> + Sync,
     {
-        let words: Result<HashMap<String, u32>> = iterator
+        let words: Result<HashMap<String, u64>> = iterator
             .maybe_par_bridge()
             .map(|sequence| {
                 let words = process(sequence.as_ref())?;
@@ -665,7 +665,7 @@ mod tests {
 
     #[test]
     fn test_train() {
-        let word_counts: HashMap<String, u32> = [
+        let word_counts: HashMap<String, u64> = [
             ("roses".into(), 1),
             ("are".into(), 2),
             ("red".into(), 1),
@@ -744,7 +744,7 @@ mod tests {
          */
 
         let max_token_length = 16;
-        let long_word_counts: HashMap<String, u32> = [
+        let long_word_counts: HashMap<String, u64> = [
             ("singlelongtokenwithoutcasechange", 2),
             ("singleLongTokenWithCamelCaseChange", 2),
             ("Longsingletokenwithpunctu@t!onwithin", 2),
@@ -784,7 +784,7 @@ mod tests {
         // directly compares tokens with known expected values.
         // maybe unstable depending on specific settings or changes.
          */
-        let long_word_counts: HashMap<String, u32> = [
+        let long_word_counts: HashMap<String, u64> = [
             ("sin", 2),
             ("Sin", 2),
             ("Lon", 2),
diff --git a/tokenizers/src/models/wordlevel/trainer.rs b/tokenizers/src/models/wordlevel/trainer.rs
index d4048b15d..c52ad08d7 100644
--- a/tokenizers/src/models/wordlevel/trainer.rs
+++ b/tokenizers/src/models/wordlevel/trainer.rs
@@ -10,7 +10,7 @@ use std::collections::HashMap;
 pub struct WordLevelTrainer {
     /// The minimum frequency a word must have to be part of the vocabulary
     #[builder(default = "0")]
-    pub min_frequency: u32,
+    pub min_frequency: u64,
     /// The target vocabulary size
     #[builder(default = "30_000")]
     pub vocab_size: usize,
@@ -22,7 +22,7 @@ pub struct WordLevelTrainer {
     pub special_tokens: Vec<AddedToken>,
 
     #[builder(default, private)]
-    words: HashMap<String, u32>,
+    words: HashMap<String, u64>,
 }
 
 impl Default for WordLevelTrainer {
@@ -38,14 +38,14 @@ impl WordLevelTrainer {
 
     fn do_train(
         &self,
-        word_counts: &HashMap<String, u32>,
+        word_counts: &HashMap<String, u64>,
         model: &mut WordLevel,
     ) -> Result<Vec<AddedToken>> {
         let mut ordered_counts = word_counts.iter().collect::<Vec<_>>();
 
         //sort the word counts first by inverse counts and then by word, in order
         //to keep the sorting deterministic in case of equal counts
-        let cmp = |l: &(&String, &u32), r: &(&String, &u32)| -> Ordering {
+        let cmp = |l: &(&String, &u64), r: &(&String, &u64)| -> Ordering {
             let count_comp: Ordering = l.1.cmp(r.1);
             if count_comp != Ordering::Equal {
                 return count_comp.reverse();
@@ -100,7 +100,7 @@ impl Trainer for WordLevelTrainer {
         S: AsRef<str> + Send,
         F: Fn(&str) -> Result<Vec<String>> + Sync,
     {
-        let words: Result<HashMap<String, u32>> = iterator
+        let words: Result<HashMap<String, u64>> = iterator
             .maybe_par_bridge()
             .map(|sequence| {
                 let words = process(sequence.as_ref())?;
@@ -132,7 +132,7 @@ mod tests {
 
     #[test]
     fn test_train() {
-        let word_counts: HashMap<String, u32> = [
+        let word_counts: HashMap<String, u64> = [
             ("the".into(), 25),
             ("roses".into(), 22),
             ("are".into(), 24),
diff --git a/tokenizers/src/models/wordpiece/trainer.rs b/tokenizers/src/models/wordpiece/trainer.rs
index 1adcc2be4..58a5abc8f 100644
--- a/tokenizers/src/models/wordpiece/trainer.rs
+++ b/tokenizers/src/models/wordpiece/trainer.rs
@@ -26,7 +26,7 @@ impl WordPieceTrainerBuilder {
 
     /// Set the expected minimum frequency
     #[must_use]
-    pub fn min_frequency(mut self, frequency: u32) -> Self {
+    pub fn min_frequency(mut self, frequency: u64) -> Self {
         self.bpe_trainer_builder = self.bpe_trainer_builder.min_frequency(frequency);
         self
     }
@@ -94,11 +94,11 @@ pub struct WordPieceTrainer {
 }
 
 impl WordPieceTrainer {
-    pub fn min_frequency(&self) -> u32 {
+    pub fn min_frequency(&self) -> u64 {
         self.bpe_trainer.min_frequency
     }
 
-    pub fn set_min_frequency(&mut self, freq: u32) {
+    pub fn set_min_frequency(&mut self, freq: u64) {
         self.bpe_trainer.min_frequency = freq;
     }
 


Test patch: