diff --git a/model2vec/model.py b/model2vec/model.py index c5f4ecc..0379736 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,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: diff --git a/tests/integration/pretrained_model_metrics.py b/tests/integration/pretrained_model_metrics.py index 13dce68..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(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, diff --git a/tests/test_model.py b/tests/test_model.py index de999c4..66c114d 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 @@ -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: