fix: add distillation duplicate guard (#350)

* feat: add distillation duplicate guard

* add normalization to added tokens

Author: stephantul

model2vec/distill/distillation.py

@@ -3,7 +3,7 @@
 import logging
 import os
 import re
-from typing import cast
+from typing import Sequence, cast
 
 from huggingface_hub.hf_api import model_info
 from skeletoken import TokenizerModel
@@ -24,7 +24,7 @@
 def distill_from_model(
     model: PreTrainedModel,
     tokenizer: PreTrainedTokenizerFast,
-    vocabulary: list[str] | None = None,
+    vocabulary: Sequence[str] | None = None,
     device: str | None = None,
     pca_dims: PCADimType = 256,
     sif_coefficient: float | None = 1e-4,

model2vec/tokenizer/tokenizer.py

@@ -2,6 +2,7 @@
 
 import logging
 import re
+from typing import Sequence
 
 from skeletoken import TokenizerModel
 
@@ -10,7 +11,7 @@
 
 def clean_and_create_vocabulary(
     model: TokenizerModel,
-    vocabulary_to_add: list[str],
+    vocabulary_to_add: Sequence[str],
     token_remove_regex: re.Pattern[str] | None,
 ) -> TokenizerModel:
     """Clean a vocabulary by removing duplicates and tokens that were already in the vocabulary.
@@ -43,6 +44,7 @@ def clean_and_create_vocabulary(
     seen_tokens = set(internal_tokens)
     tokens_to_add: list[str] = []
     added_tokens_to_add: list[str] = []
+    seen_added = set()
     for token in vocabulary_to_add:
         preprocessed = preprocessor.preprocess(token)
         if len(preprocessed) < 1:
@@ -58,7 +60,15 @@ def clean_and_create_vocabulary(
             if token in model.vocabulary:
                 # If the unprocessed token (incorrectly) is in the vocabulary, we should remove it.
                 model = model.remove_token_from_vocabulary(token)
+            if preprocessor.normalizer:
+                token = preprocessor.normalizer.normalize_str(token)
+            # We need to strip because our AddedTokens also get stripped
+            token = token.strip()
+            if token in seen_added:
+                logger.warning(f"Normalized added token '{token}' was in the added vocabulary twice.")
+                continue
             added_tokens_to_add.append(token)
+            seen_added.add(token)
             continue
         token = preprocessed[0]
         if token in seen_tokens:

tests/test_distillation.py

@@ -319,6 +319,10 @@ def test__post_process_embeddings(
         # Case: Only empty token (""), should return an empty list
         ([""], [], ["empty"]),
         (["multi word token"], ["multi word token"], []),
+        # Case: duplicate multi-word token, should only be added once
+        (["multi word token", "multi word token"], ["multi word token"], ["twice"]),
+        # Case: multi-word tokens that normalize to the same string, should only be added once
+        (["Multi Word Token", "multi word token"], ["multi word token"], ["twice"]),
     ],
 )
 def test_clean_and_create_vocabulary(