fix: Incorrect added token can cause issues when adding token as multiword token

model2vec/distill/distillation.py

@@ -9,9 +9,8 @@
 from huggingface_hub.hf_api import model_info
 from skeletoken import TokenizerModel
 from skeletoken.external.transformers import reshape_embeddings
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer, PreTrainedTokenizerFast
 from transformers.modeling_utils import PreTrainedModel
-from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 
 from model2vec.distill.inference import PCADimType, PoolingMode, create_embeddings, post_process_embeddings
 from model2vec.distill.utils import select_optimal_device

model2vec/tokenizer/tokenizer.py

@@ -53,6 +53,9 @@ def clean_and_create_vocabulary(
             logger.warning(
                 f"Token '{token}' was split into multiple tokens after preprocessing: [{split_into}], adding it as a multi-word token."
             )
+            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)
             added_tokens_to_add.append(token)
             continue
         token = preprocessed[0]

tests/test_distillation.py

@@ -9,9 +9,9 @@
 import pytest
 from pytest import LogCaptureFixture
 from skeletoken import TokenizerModel
-from transformers import BertTokenizerFast
+from transformers import BertTokenizer
 from transformers.modeling_utils import PreTrainedModel
-from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+from transformers.tokenization_utils_tokenizers import PreTrainedTokenizerFast
 
 from model2vec.distill.distillation import distill, distill_from_model
 from model2vec.distill.inference import PoolingMode, create_embeddings, post_process_embeddings
@@ -38,6 +38,7 @@
         (None, None, 1e-4),  # No PCA, SIF on
         (None, 0.9, 1e-4),  # PCA as float (variance), SIF on
         (["star wars"], 8, None),  # Multiword vocabulary
+        (["..."], 8, None),  # Crashing multiword vocabulary
     ],
 )
 @patch.object(import_module("model2vec.distill.distillation"), "model_info")
@@ -92,7 +93,7 @@ def test_distill_from_model(
 def test_distill_removal_pattern_all_tokens(
     mock_auto_model: MagicMock,
     mock_model_info: MagicMock,
-    mock_berttokenizer: BertTokenizerFast,
+    mock_berttokenizer: BertTokenizer,
     mock_transformer: PreTrainedModel,
 ) -> None:
     """Test the removal pattern."""