feat: max length behavior

model2vec/model.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import json
 import math
 import os
@@ -77,7 +78,7 @@ 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]))
@@ -129,6 +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
+        self._set_max_length_in_tokenizer(value)
 
     @property
     def embedding_dtype(self) -> str:
@@ -163,18 +165,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 +180,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,10 +231,18 @@ 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,
         batch_size: int = 1024,
         show_progress_bar: bool = False,
@@ -252,6 +254,7 @@ def encode_as_sequence(
     def encode_as_sequence(
         self,
         sentences: list[str],
+        *,
         max_length: int | None = None,
         batch_size: int = 1024,
         show_progress_bar: bool = False,
@@ -262,6 +265,7 @@ def encode_as_sequence(
     def encode_as_sequence(
         self,
         sentences: str | list[str],
+        *,
         max_length: int | None = None,
         batch_size: int = 1024,
         show_progress_bar: bool = False,
@@ -275,13 +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 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.
+            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,38 +299,45 @@ 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)
 
-        # 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))
+        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:
+            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:
@@ -375,30 +388,37 @@ 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)
 
-        # 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)
+        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, 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, normalize))
+                out_array = np.concatenate(out_arrays, axis=0)
+        finally:
+            self._set_max_length_in_tokenizer(self.max_length)
 
         if was_single:
             return out_array[0]
@@ -424,9 +444,9 @@ 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, max_length=max_length)
+        ids = self.tokenize(sentences=sentences)
         out: list[np.ndarray] = []
         for id_list in ids:
             if id_list:

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": model.median_token_length,
         "unk_token_id": model.unk_token_id,
         "normalize": bool(model.normalize),
         "base_model_name": model.base_model_name,

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
 
@@ -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:
@@ -272,9 +284,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
 
@@ -301,8 +310,24 @@ 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:
+    """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: