fix(train): drop unknown tokens during training and inference

model2vec/train/base.py

@@ -12,7 +12,7 @@
 from tqdm import trange
 
 from model2vec.inference import StaticModelPipeline
-from model2vec.model import DEFAULT_MAX_LENGTH, PathLike, StaticModel
+from model2vec.model import DEFAULT_MAX_LENGTH, PathLike, StaticModel, _get_unk_token_id
 from model2vec.train.dataset import TextDataset
 from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop
 from model2vec.train.utils import (
@@ -90,6 +90,13 @@ def __init__(
         self._weights = weights
         self.w = self.construct_weights()
         self.tokenizer = tokenizer
+        self.unk_token_id = _get_unk_token_id(tokenizer)
+
+    def _remove_unk(self, token_ids: list[int]) -> list[int]:
+        """Drop unknown tokens, mirroring `StaticModel.tokenize`."""
+        if self.unk_token_id is None:
+            return token_ids
+        return [token_id for token_id in token_ids if token_id != self.unk_token_id]
 
     def construct_weights(self) -> nn.Parameter:
         """Construct the weights for the model."""
@@ -248,7 +255,9 @@ def tokenize(self, texts: list[str]) -> torch.Tensor:
         """
         max_length = self.max_length
         encoded: list[Encoding] = self.tokenizer.encode_batch_fast(texts, add_special_tokens=False)
-        encoded_ids: list[torch.Tensor] = [torch.Tensor(encoding.ids[:max_length]).long() for encoding in encoded]
+        encoded_ids: list[torch.Tensor] = [
+            torch.Tensor(self._remove_unk(encoding.ids)[:max_length]).long() for encoding in encoded
+        ]
         return pad_sequence(encoded_ids, batch_first=True, padding_value=self.pad_id)
 
     @property
@@ -400,7 +409,7 @@ def _prepare_dataset(self, X: list[str], y: torch.Tensor, max_length: int | None
                 truncate_length = max_length * 10
                 batch = [x[:truncate_length] for x in batch]
             encoded = self.tokenizer.encode_batch_fast(batch, add_special_tokens=False)
-            tokenized.extend([encoding.ids[:max_length] for encoding in encoded])
+            tokenized.extend([self._remove_unk(encoding.ids)[:max_length] for encoding in encoded])
 
         return TextDataset(tokenized, y, pad_id=self.pad_id)
 

tests/test_trainable.py

@@ -160,6 +160,17 @@ def test_training_batch_padding_is_masked(mock_vectors: np.ndarray, mock_tokeniz
         assert torch.allclose(s._encode(batch)[0], s._encode(s.tokenize(texts[:1]))[0])
 
 
+def test_unknown_tokens_are_dropped(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None:
+    """Training and inference should drop unknown tokens, the way `StaticModel.tokenize` does."""
+    s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer)
+    static = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer)
+    texts = ["word1 unknownword", "unknownword word2 otherunknown"]
+    expected = static.tokenize(texts)
+
+    assert [row[row != s.pad_id].tolist() for row in s.tokenize(texts)] == expected
+    assert s._prepare_dataset(texts, torch.arange(2), max_length=None).tokenized_texts == expected
+
+
 def test_predict(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Test the predict function."""
     result = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist()