From e0fcdfaf34fbe74beca3635fcae845bd820389d6 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 12:35:44 +0200 Subject: [PATCH 1/8] fix: set max length behavior --- model2vec/model.py | 128 ++++++++++-------- tests/integration/pretrained_model_metrics.py | 2 +- .../test_pretrained_model_regression.py | 3 +- tests/test_model.py | 16 ++- 4 files changed, 88 insertions(+), 61 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index c5f4ecc..596123a 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy import json import math import os @@ -77,10 +78,9 @@ def __init__( # We can't use or short circuit here because np.ndarray as booleans are ambiguous. self.token_mapping: np.ndarray | None = token_mapping - self.tokenizer = tokenizer + self.tokenizer = copy.deepcopy(tokenizer) self.unk_token_id = _get_unk_token_id(self.tokenizer) - self.median_token_length = int(np.median([len(token) for token in self.tokens])) self.config: StaticModelConfig = cast(StaticModelConfig, {**config}) if config is not None else {} self.base_model_name = base_model_name self.language = language @@ -129,6 +129,10 @@ def max_length(self, value: int | None) -> None: f"Set max_length to `{value}`, which does not match config value `{config_max_length}`. Updating config." ) self.config["max_length"] = value + if value is None: + self.tokenizer.no_truncation() + else: + self.tokenizer.enable_truncation(value) @property def embedding_dtype(self) -> str: @@ -163,18 +167,12 @@ def save_pretrained(self, path: PathLike, model_name: str | None = None, subfold mapping=self.token_mapping, ) - def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> list[list[int]]: + def tokenize(self, sentences: Sequence[str]) -> list[list[int]]: """Tokenize a list of sentences. :param sentences: The sentences to tokenize. - :param max_length: The maximum length of the sentences in tokens. If this is None, sequences - are not truncated. :return: A list of list of tokens. """ - if max_length is not None: - m = max_length * self.median_token_length - sentences = [sentence[:m] for sentence in sentences] - encodings: list[Encoding] = self.tokenizer.encode_batch_fast(sentences, add_special_tokens=False) encodings_ids = [encoding.ids for encoding in encodings] @@ -184,8 +182,6 @@ def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> l encodings_ids = [ [token_id for token_id in token_ids if token_id != self.unk_token_id] for token_ids in encodings_ids ] - if max_length is not None: - encodings_ids = [token_ids[:max_length] for token_ids in encodings_ids] return encodings_ids @@ -237,11 +233,19 @@ def from_pretrained( force_download=force_download, ) + def _set_max_length_in_tokenizer(self, max_length: int | None) -> None: + """Sets the max length in the tokenizer.""" + if max_length is None: + self.tokenizer.no_truncation() + else: + self.tokenizer.enable_truncation(max_length) + @overload def encode_as_sequence( self, sentences: str, - max_length: int | None = None, + *, + max_length: int | None | _UnsetType = _UNSET, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -252,7 +256,8 @@ def encode_as_sequence( def encode_as_sequence( self, sentences: list[str], - max_length: int | None = None, + *, + max_length: int | None | _UnsetType = _UNSET, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -262,7 +267,8 @@ def encode_as_sequence( def encode_as_sequence( self, sentences: str | list[str], - max_length: int | None = None, + *, + max_length: int | None | _UnsetType = _UNSET, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -298,33 +304,39 @@ def encode_as_sequence( sentence_batches = list(self._batch(sentences, batch_size)) total_batches = math.ceil(len(sentences) / batch_size) - # Use joblib for multiprocessing if requested, and if we have enough sentences - if use_multiprocessing and len(sentences) > multiprocessing_threshold: - # Disable parallelism for tokenizers - os.environ["TOKENIZERS_PARALLELISM"] = "false" - - results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)( - delayed(self._encode_batch_as_sequence)(batch, max_length) for batch in sentence_batches - ) - out_array: list[np.ndarray] = [] - for r in results: - out_array.extend(r) - else: - out_array = [] - for batch in tqdm( - sentence_batches, - total=total_batches, - disable=not show_progress_bar, - ): - out_array.extend(self._encode_batch_as_sequence(batch, max_length)) + if not isinstance(max_length, _UnsetType): + self._set_max_length_in_tokenizer(max_length) + try: + # Use joblib for multiprocessing if requested, and if we have enough sentences + if use_multiprocessing and len(sentences) > multiprocessing_threshold: + # Disable parallelism for tokenizers + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)( + delayed(self._encode_batch_as_sequence)(batch) for batch in sentence_batches + ) + out_array: list[np.ndarray] = [] + for r in results: + out_array.extend(r) + else: + out_array = [] + for batch in tqdm( + sentence_batches, + total=total_batches, + disable=not show_progress_bar, + ): + out_array.extend(self._encode_batch_as_sequence(batch)) + finally: + if not isinstance(max_length, _UnsetType): + self._set_max_length_in_tokenizer(self.max_length) if was_single: return out_array[0] return out_array - def _encode_batch_as_sequence(self, sentences: Sequence[str], max_length: int | None) -> list[np.ndarray]: + def _encode_batch_as_sequence(self, sentences: Sequence[str]) -> list[np.ndarray]: """Encode a batch of sentences as a sequence.""" - ids = self.tokenize(sentences=sentences, max_length=max_length) + ids = self.tokenize(sentences=sentences) out: list[np.ndarray] = [] for id_list in ids: if id_list: @@ -380,25 +392,31 @@ def encode( sentence_batches = list(self._batch(sentences, batch_size)) total_batches = math.ceil(len(sentences) / batch_size) - # Use joblib for multiprocessing if requested, and if we have enough sentences - if use_multiprocessing and len(sentences) > multiprocessing_threshold: - # Disable parallelism for tokenizers - os.environ["TOKENIZERS_PARALLELISM"] = "false" - - results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)( - delayed(self._encode_batch)(batch, max_length, normalize) for batch in sentence_batches - ) - out_array = np.concatenate(results, axis=0) - else: - # Don't use multiprocessing - out_arrays: list[np.ndarray] = [] - for batch in tqdm( - sentence_batches, - total=total_batches, - disable=not show_progress_bar, - ): - out_arrays.append(self._encode_batch(batch, max_length, normalize)) - out_array = np.concatenate(out_arrays, axis=0) + if not isinstance(max_length, _UnsetType): + self._set_max_length_in_tokenizer(max_length) + try: + # Use joblib for multiprocessing if requested, and if we have enough sentences + if use_multiprocessing and len(sentences) > multiprocessing_threshold: + # Disable parallelism for tokenizers + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)( + delayed(self._encode_batch)(batch, max_length, normalize) for batch in sentence_batches + ) + out_array = np.concatenate(results, axis=0) + else: + # Don't use multiprocessing + out_arrays: list[np.ndarray] = [] + for batch in tqdm( + sentence_batches, + total=total_batches, + disable=not show_progress_bar, + ): + out_arrays.append(self._encode_batch(batch, max_length, normalize)) + out_array = np.concatenate(out_arrays, axis=0) + finally: + if not isinstance(max_length, _UnsetType): + self._set_max_length_in_tokenizer(self.max_length) if was_single: return out_array[0] @@ -426,7 +444,7 @@ def _encode_helper(self, id_list: list[int]) -> np.ndarray: def _encode_batch(self, sentences: Sequence[str], max_length: int | None, normalize: bool) -> np.ndarray: """Encode a batch of sentences.""" - ids = self.tokenize(sentences=sentences, max_length=max_length) + ids = self.tokenize(sentences=sentences) out: list[np.ndarray] = [] for id_list in ids: if id_list: diff --git a/tests/integration/pretrained_model_metrics.py b/tests/integration/pretrained_model_metrics.py index 13dce68..39450e4 100644 --- a/tests/integration/pretrained_model_metrics.py +++ b/tests/integration/pretrained_model_metrics.py @@ -87,7 +87,7 @@ def compute_metrics(model: StaticModel) -> dict[str, Any]: "token_order_hash": token_order_hash, "first_tokens": tokens[:10], "last_tokens": tokens[-10:], - "median_token_length": int(model.median_token_length), + "median_token_length": int(np.median([len(token) for token in tokens])), "unk_token_id": model.unk_token_id, "normalize": bool(model.normalize), "base_model_name": model.base_model_name, diff --git a/tests/integration/test_pretrained_model_regression.py b/tests/integration/test_pretrained_model_regression.py index 21d5b84..dd5e437 100644 --- a/tests/integration/test_pretrained_model_regression.py +++ b/tests/integration/test_pretrained_model_regression.py @@ -77,7 +77,8 @@ def test_all_attributes_are_loaded(model: StaticModel) -> None: assert model.embedding_dtype == np.dtype(model.embedding.dtype).name assert isinstance(model.config, dict) and model.config assert isinstance(model.normalize, bool) - assert isinstance(model.median_token_length, int) and model.median_token_length > 0 + median_token_length = np.median([len(token) for token in model.tokens]) + assert median_token_length > 0 assert model.unk_token_id is None or isinstance(model.unk_token_id, int) assert model.base_model_name is None or isinstance(model.base_model_name, str) assert model.language is None or isinstance(model.language, list) diff --git a/tests/test_model.py b/tests/test_model.py index de999c4..08b70b0 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -13,7 +13,7 @@ def test_initialization(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, moc model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config) assert model.embedding.shape == (5, 2) assert len(model.tokens) == 5 - assert model.tokenizer == mock_tokenizer + assert model.tokenizer.get_vocab() == mock_tokenizer.get_vocab() assert model.config == {**original_config, "normalize": False, "max_length": 512} assert mock_config == original_config @@ -272,9 +272,6 @@ def test_load_pretrained_vocabulary_quantized( def test_initialize_normalize(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """Tests whether the normalization initialization is correct.""" - model = StaticModel(mock_vectors, mock_tokenizer, {}, normalize=None) - assert not model.normalize - model = StaticModel(mock_vectors, mock_tokenizer, {}, normalize=False) assert not model.normalize @@ -305,6 +302,17 @@ def test_set_max_length(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> assert model.config == {"normalize": False, "max_length": 256} +def test_models_do_not_share_tokenizer(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Tests that two models built from the same tokenizer object have independent truncation state.""" + model_a = StaticModel(mock_vectors, mock_tokenizer, {}, max_length=128) + model_b = StaticModel(mock_vectors, mock_tokenizer, {}, max_length=256) + + model_b.max_length = 8 + + assert model_a.tokenizer.truncation["max_length"] == 128 + assert model_b.tokenizer.truncation["max_length"] == 8 + + def test_dim(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]) -> None: """Tests the dimensionality of the model.""" model = StaticModel(mock_vectors, mock_tokenizer, mock_config) From 06cc82deafee1a18cd6a6af170323b43fc305601 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 12:50:17 +0200 Subject: [PATCH 2/8] restore max length --- model2vec/model.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 596123a..c8b02f0 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -81,6 +81,7 @@ def __init__( self.tokenizer = copy.deepcopy(tokenizer) self.unk_token_id = _get_unk_token_id(self.tokenizer) + self.median_token_length = int(np.median([len(token) for token in self.tokens])) self.config: StaticModelConfig = cast(StaticModelConfig, {**config}) if config is not None else {} self.base_model_name = base_model_name self.language = language @@ -299,13 +300,17 @@ def encode_as_sequence( if isinstance(sentences, str): sentences = [sentences] was_single = True + if isinstance(max_length, _UnsetType): + max_length = self.max_length + if max_length is not None: + m = max_length * self.median_token_length + sentences = [sentence[:m] for sentence in sentences] # Prepare all batches sentence_batches = list(self._batch(sentences, batch_size)) total_batches = math.ceil(len(sentences) / batch_size) - if not isinstance(max_length, _UnsetType): - self._set_max_length_in_tokenizer(max_length) + self._set_max_length_in_tokenizer(max_length) try: # Use joblib for multiprocessing if requested, and if we have enough sentences if use_multiprocessing and len(sentences) > multiprocessing_threshold: @@ -327,8 +332,7 @@ def encode_as_sequence( ): out_array.extend(self._encode_batch_as_sequence(batch)) finally: - if not isinstance(max_length, _UnsetType): - self._set_max_length_in_tokenizer(self.max_length) + self._set_max_length_in_tokenizer(self.max_length) if was_single: return out_array[0] @@ -387,13 +391,15 @@ def encode( max_length = self.max_length if normalize is None: normalize = self.normalize + if max_length is not None: + m = max_length * self.median_token_length + sentences = [sentence[:m] for sentence in sentences] # Prepare all batches sentence_batches = list(self._batch(sentences, batch_size)) total_batches = math.ceil(len(sentences) / batch_size) - if not isinstance(max_length, _UnsetType): - self._set_max_length_in_tokenizer(max_length) + self._set_max_length_in_tokenizer(max_length) try: # Use joblib for multiprocessing if requested, and if we have enough sentences if use_multiprocessing and len(sentences) > multiprocessing_threshold: @@ -401,7 +407,7 @@ def encode( os.environ["TOKENIZERS_PARALLELISM"] = "false" results = ProgressParallel(n_jobs=-1, use_tqdm=show_progress_bar, total=total_batches)( - delayed(self._encode_batch)(batch, max_length, normalize) for batch in sentence_batches + delayed(self._encode_batch)(batch, normalize) for batch in sentence_batches ) out_array = np.concatenate(results, axis=0) else: @@ -412,11 +418,10 @@ def encode( total=total_batches, disable=not show_progress_bar, ): - out_arrays.append(self._encode_batch(batch, max_length, normalize)) + out_arrays.append(self._encode_batch(batch, normalize)) out_array = np.concatenate(out_arrays, axis=0) finally: - if not isinstance(max_length, _UnsetType): - self._set_max_length_in_tokenizer(self.max_length) + self._set_max_length_in_tokenizer(self.max_length) if was_single: return out_array[0] @@ -442,7 +447,7 @@ def _encode_helper(self, id_list: list[int]) -> np.ndarray: return emb - def _encode_batch(self, sentences: Sequence[str], max_length: int | None, normalize: bool) -> np.ndarray: + def _encode_batch(self, sentences: Sequence[str], normalize: bool) -> np.ndarray: """Encode a batch of sentences.""" ids = self.tokenize(sentences=sentences) out: list[np.ndarray] = [] From 8f6a9c591c0171864d187841637ca3e6530152d9 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 12:51:54 +0200 Subject: [PATCH 3/8] consolidate function --- model2vec/model.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index c8b02f0..3e399d3 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -130,10 +130,7 @@ def max_length(self, value: int | None) -> None: f"Set max_length to `{value}`, which does not match config value `{config_max_length}`. Updating config." ) self.config["max_length"] = value - if value is None: - self.tokenizer.no_truncation() - else: - self.tokenizer.enable_truncation(value) + self._set_max_length_in_tokenizer(value) @property def embedding_dtype(self) -> str: From 613bc0476ea01e341dc5e8ad651f1917f3986b0b Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 12:55:11 +0200 Subject: [PATCH 4/8] revert median behavior in integration test --- tests/integration/pretrained_model_metrics.py | 2 +- tests/integration/test_pretrained_model_regression.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/integration/pretrained_model_metrics.py b/tests/integration/pretrained_model_metrics.py index 39450e4..66a5e76 100644 --- a/tests/integration/pretrained_model_metrics.py +++ b/tests/integration/pretrained_model_metrics.py @@ -87,7 +87,7 @@ def compute_metrics(model: StaticModel) -> dict[str, Any]: "token_order_hash": token_order_hash, "first_tokens": tokens[:10], "last_tokens": tokens[-10:], - "median_token_length": int(np.median([len(token) for token in tokens])), + "median_token_length": model.median_token_length, "unk_token_id": model.unk_token_id, "normalize": bool(model.normalize), "base_model_name": model.base_model_name, diff --git a/tests/integration/test_pretrained_model_regression.py b/tests/integration/test_pretrained_model_regression.py index dd5e437..21d5b84 100644 --- a/tests/integration/test_pretrained_model_regression.py +++ b/tests/integration/test_pretrained_model_regression.py @@ -77,8 +77,7 @@ def test_all_attributes_are_loaded(model: StaticModel) -> None: assert model.embedding_dtype == np.dtype(model.embedding.dtype).name assert isinstance(model.config, dict) and model.config assert isinstance(model.normalize, bool) - median_token_length = np.median([len(token) for token in model.tokens]) - assert median_token_length > 0 + assert isinstance(model.median_token_length, int) and model.median_token_length > 0 assert model.unk_token_id is None or isinstance(model.unk_token_id, int) assert model.base_model_name is None or isinstance(model.base_model_name, str) assert model.language is None or isinstance(model.language, list) From 5663b449a5dcf4837b03f93b68cef0a97a4f431c Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 13:02:12 +0200 Subject: [PATCH 5/8] add test --- tests/test_model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_model.py b/tests/test_model.py index 08b70b0..f8d7a78 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -298,8 +298,13 @@ def test_set_max_length(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> """Tests whether the max_length is set correctly.""" model = StaticModel(mock_vectors, mock_tokenizer, {}, max_length=128) assert model.config == {"normalize": False, "max_length": 128} + assert model.tokenizer.truncation["max_length"] == 128 model.max_length = 256 assert model.config == {"normalize": False, "max_length": 256} + assert model.tokenizer.truncation["max_length"] == 256 + model.max_length = None + assert model.config == {"normalize": False, "max_length": None} + assert model.tokenizer.truncation is None def test_models_do_not_share_tokenizer(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: From e3215f0cc1da0affa11d57ba184be256961806f9 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 14:13:48 +0200 Subject: [PATCH 6/8] remove max length from encode_as_sequence --- model2vec/model.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 3e399d3..47f9f44 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -243,7 +243,6 @@ def encode_as_sequence( self, sentences: str, *, - max_length: int | None | _UnsetType = _UNSET, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -255,7 +254,6 @@ def encode_as_sequence( self, sentences: list[str], *, - max_length: int | None | _UnsetType = _UNSET, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -266,7 +264,6 @@ def encode_as_sequence( self, sentences: str | list[str], *, - max_length: int | None | _UnsetType = _UNSET, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -279,13 +276,12 @@ def encode_as_sequence( Note that if you just want the mean, you should use the `encode` method. This is about twice as slow. Sentences that do not contain any tokens will be turned into an empty array. + Unlike `encode`, this method never truncates: the full sequence of token embeddings is always returned. NOTE: the input type is currently underspecified. The actual input type is `Sequence[str] | str`, but this is not possible to implement in python typing currently. :param sentences: The list of sentences to encode. - :param max_length: The maximum length of the sentences. Any tokens beyond this length will be truncated. - If this is None, no truncation is done. :param batch_size: The batch size to use. :param show_progress_bar: Whether to show the progress bar. :param use_multiprocessing: Whether to use multiprocessing. @@ -297,17 +293,12 @@ def encode_as_sequence( if isinstance(sentences, str): sentences = [sentences] was_single = True - if isinstance(max_length, _UnsetType): - max_length = self.max_length - if max_length is not None: - m = max_length * self.median_token_length - sentences = [sentence[:m] for sentence in sentences] # Prepare all batches sentence_batches = list(self._batch(sentences, batch_size)) total_batches = math.ceil(len(sentences) / batch_size) - self._set_max_length_in_tokenizer(max_length) + self._set_max_length_in_tokenizer(None) try: # Use joblib for multiprocessing if requested, and if we have enough sentences if use_multiprocessing and len(sentences) > multiprocessing_threshold: From 58e00c76af179a27add1f39f46c099613d4898d1 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 15:01:06 +0200 Subject: [PATCH 7/8] set max length for encode_as_sequence --- model2vec/model.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 47f9f44..0379736 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -243,6 +243,7 @@ def encode_as_sequence( self, sentences: str, *, + max_length: int | None = None, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -254,6 +255,7 @@ def encode_as_sequence( self, sentences: list[str], *, + max_length: int | None = None, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -264,6 +266,7 @@ def encode_as_sequence( self, sentences: str | list[str], *, + max_length: int | None = None, batch_size: int = 1024, show_progress_bar: bool = False, use_multiprocessing: bool = True, @@ -276,12 +279,15 @@ def encode_as_sequence( Note that if you just want the mean, you should use the `encode` method. This is about twice as slow. Sentences that do not contain any tokens will be turned into an empty array. - Unlike `encode`, this method never truncates: the full sequence of token embeddings is always returned. + Unlike `encode`, this method ignores the model's `max_length`: truncation is only applied if you + explicitly pass `max_length` here. NOTE: the input type is currently underspecified. The actual input type is `Sequence[str] | str`, but this is not possible to implement in python typing currently. :param sentences: The list of sentences to encode. + :param max_length: The maximum length of the sentences. Any tokens beyond this length will be truncated. + If this is None, no truncation is done. Note that this is independent of the model's `max_length`. :param batch_size: The batch size to use. :param show_progress_bar: Whether to show the progress bar. :param use_multiprocessing: Whether to use multiprocessing. @@ -293,12 +299,15 @@ def encode_as_sequence( if isinstance(sentences, str): sentences = [sentences] was_single = True + if max_length is not None: + m = max_length * self.median_token_length + sentences = [sentence[:m] for sentence in sentences] # Prepare all batches sentence_batches = list(self._batch(sentences, batch_size)) total_batches = math.ceil(len(sentences) / batch_size) - self._set_max_length_in_tokenizer(None) + self._set_max_length_in_tokenizer(max_length) try: # Use joblib for multiprocessing if requested, and if we have enough sentences if use_multiprocessing and len(sentences) > multiprocessing_threshold: From 07c734f46491ecba7c17dfac47cf8e3f580580de Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 11 Sep 2026 15:31:15 +0200 Subject: [PATCH 8/8] update tests --- tests/test_model.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_model.py b/tests/test_model.py index f8d7a78..66c114d 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -86,6 +86,18 @@ def test_encode_multiprocessing( assert encoded.shape == (15000, 2) +def test_encode_as_sequence_max_length( + mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str] +) -> None: + """Test encoding of sentences as tokens with max_length truncation.""" + sentences = ["word1 word2 word3"] + model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config) + full = model.encode_as_sequence(sentences) + truncated = model.encode_as_sequence(sentences, max_length=1) + + assert len(truncated[0]) < len(full[0]) + + def test_encode_as_sequence_multiprocessing( mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str] ) -> None: