Repository: huggingface/tokenizers (language: rust)

Issue:
UnigramTrainer Stalls on Small Vocabulary Size
I'm training a Unigram model using the UnigramTrainer and inexplicably, it seems to not terminate sometimes when given a large enough coprus, but not for other slightly smaller corpus sizes.

I'm training tokenizers with different vocabulary sizes to compare models from a vocab size of 100 tokens to thousands of tokens and I suspect that it has to do with a while condition not being met somewhere for small vocabulary sizes. Perhaps the initial alphabet is larger than the vocabulary size specified and this causes the problem?

I had a similar problem when I implemented my own tokenzier. If you don't check that the initial alphabet is the same size as the pruned vocabulary size, you can keep calculating loss and trying to prune non-alphabet tokens, but there are none and so you can never reach your target vocabulary size.


Discussion / hints:
Could you give a reproducible example ?

If your theory is correct it should be reproducible with "abc" and asking for `vocab_size=2` , right ?
Here is a somewhat minimal example that corresponds to my training script, just incase my training options are problematic:

```
from tokenizers import Tokenizer
from tokenizers.models import Unigram
from tokenizers.normalizers import Lowercase
from tokenizers.pre_tokenizers import WhitespaceSplit
from tokenizers.trainers import UnigramTrainer

text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna " \
"aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. " \
"Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint " \
"occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Sed ut perspiciatis " \
"unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab " \
"illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia " \
"voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi " \
"nesciunt. Neque porro quisquam est, qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit, sed quia non " \
"numquam eius modi tempora incidunt ut labore et dolore magnam aliquam quaerat voluptatem. Ut enim ad minima veniam, " \
"quis nostrum exercitationem ullam corporis suscipit laboriosam, nisi ut aliquid ex ea commodi consequatur? Quis autem " \
"vel eum iure reprehenderit qui in ea voluptate velit esse quam nihil molestiae consequatur, vel illum qui dolorem eum " \
"fugiat quo voluptas nulla pariatur?"


def get_chunks(text, maxlength):

    while text:
        yield text[:maxlength]
        text = text[maxlength:]


corpus = list(get_chunks(text, 20))

alphabet_size = len(set([char for sent in corpus for char in sent]))

for vocab_size in range(500, 1, -1):
    tokenizer = Tokenizer(Unigram())
    tokenizer.normalizer = Lowercase()
    tokenizer.pre_tokenizer = WhitespaceSplit()
    trainer = UnigramTrainer(vocab_size=vocab_size,
                             shrinking_factor=0.7)

    tokenizer.train_from_iterator(corpus, trainer)
    print(alphabet_size, vocab_size)
```
    
 This works until the vocab size is reduced to 21 on my machine before it hangs.

Fix patch:
diff --git a/tokenizers/src/models/unigram/trainer.rs b/tokenizers/src/models/unigram/trainer.rs
index d4562b28b..ab1bb6747 100644
--- a/tokenizers/src/models/unigram/trainer.rs
+++ b/tokenizers/src/models/unigram/trainer.rs
@@ -29,6 +29,12 @@ fn digamma(mut x: f64) -> f64 {
     result
 }
 
+#[derive(thiserror::Error, Debug)]
+pub enum UnigramTrainerError {
+    #[error("The vocabulary is not large enough to contain all chars")]
+    VocabularyTooSmall,
+}
+
 fn to_log_prob(pieces: &mut [SentencePiece]) {
     let sum: f64 = pieces.iter().map(|(_, score)| score).sum();
     let logsum = sum.ln();
@@ -504,6 +510,9 @@ impl UnigramTrainer {
         let expected_updates = expected_loops * self.n_sub_iterations as usize;
         self.update_progress(&progress, expected_updates, "EM training");
         let required_chars = self.required_chars(&sentences);
+        if required_chars.len() as u32 > self.vocab_size {
+            return Err(Box::new(UnigramTrainerError::VocabularyTooSmall));
+        }
         let mut new_model = Unigram::from(pieces.clone(), Some(0))?;
         loop {
             // Sub-EM iteration.
diff --git a/tokenizers/src/utils/truncation.rs b/tokenizers/src/utils/truncation.rs
index 648147872..a8ad2a614 100644
--- a/tokenizers/src/utils/truncation.rs
+++ b/tokenizers/src/utils/truncation.rs
@@ -3,16 +3,12 @@ use serde::{Deserialize, Serialize};
 use std::cmp;
 use std::mem;
 
-#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Eq, Default)]
 pub enum TruncationDirection {
     Left,
+    #[default]
     Right,
 }
-impl Default for TruncationDirection {
-    fn default() -> Self {
-        TruncationDirection::Right
-    }
-}
 
 impl std::convert::AsRef<str> for TruncationDirection {
     fn as_ref(&self) -> &str {


Test patch: