From 4976cc8aea49d2dd603853f01fd9894faabb0d30 Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Thu, 3 Sep 2026 15:44:51 +0700 Subject: [PATCH 1/3] fix(train): drop unknown tokens during training and inference --- model2vec/train/base.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/model2vec/train/base.py b/model2vec/train/base.py index 8f6b4f2..a60a150 100644 --- a/model2vec/train/base.py +++ b/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) From 97beaae4dfc1e2e200a0aaaecda6ad201c549a81 Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Thu, 3 Sep 2026 15:55:09 +0700 Subject: [PATCH 2/3] test: cover unk stripping against StaticModel tokenization --- tests/test_trainable.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 9bb4c63..fd74a1f 100644 --- a/tests/test_trainable.py +++ b/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() From 357b101477d17208afdfb596e52bf9cd5894e9aa Mon Sep 17 00:00:00 2001 From: serhiizghama Date: Fri, 4 Sep 2026 15:43:49 +0700 Subject: [PATCH 3/3] test: cover the no-unk-token path in train tokenization --- tests/test_trainable.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_trainable.py b/tests/test_trainable.py index fd74a1f..f8b2af2 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -6,6 +6,8 @@ import torch from skeletoken import TokenizerModel from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.pre_tokenizers import Whitespace from torch import nn from torch.utils.data import DataLoader, TensorDataset from transformers import AutoTokenizer @@ -171,6 +173,24 @@ def test_unknown_tokens_are_dropped(mock_vectors: np.ndarray, mock_tokenizer: To assert s._prepare_dataset(texts, torch.arange(2), max_length=None).tokenized_texts == expected +def test_tokenize_without_unk_token(mock_vectors: np.ndarray) -> None: + """A tokenizer with no unk token has nothing to drop, so tokenization is left as is.""" + vocab = ["[PAD]", "word1", "word2", "word3"] + tokenizer = Tokenizer(BPE(vocab={token: idx for idx, token in enumerate(vocab)}, merges=[], ignore_merges=True)) + tokenizer.pre_tokenizer = Whitespace() # type: ignore[assignment] + vectors = mock_vectors[: len(vocab)] + + s = StaticModelForClassification(vectors=torch.from_numpy(vectors).float(), tokenizer=tokenizer) + static = StaticModel(vectors=vectors, tokenizer=tokenizer) + assert s.unk_token_id is None + + texts = ["word1 word2", "word3 word1 word2"] + 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()