From 156a522bb7136aedf3e37643a5c1616356546494 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 12 Sep 2026 08:37:40 +0200 Subject: [PATCH] feat: small helper method to reduce overlap --- model2vec/model.py | 146 ++++++++++++++++++++++++--------------------- 1 file changed, 77 insertions(+), 69 deletions(-) diff --git a/model2vec/model.py b/model2vec/model.py index 0379736..6c7dde6 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -238,6 +238,60 @@ def _set_max_length_in_tokenizer(self, max_length: int | None) -> None: else: self.tokenizer.enable_truncation(max_length) + def _encode_dispatch( + self, + sentences: str | Sequence[str], + *, + max_length: int | None, + batch_size: int, + show_progress_bar: bool, + use_multiprocessing: bool, + multiprocessing_threshold: int, + batch_fn: Any, + batch_args: tuple[Any, ...] = (), + ) -> tuple[list[Any], bool]: + """Shared truncation, batching, and multiprocessing dispatch logic for `encode` and `encode_as_sequence`. + + :param sentences: The sentence or sentences to encode. + :param max_length: The maximum length of the sentences. Any tokens beyond this length are truncated. + :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. + :param multiprocessing_threshold: The threshold in number of sentences for using multiprocessing. + :param batch_fn: The function to apply to each batch of sentences. + :param batch_args: Additional positional arguments passed to `batch_fn` after the batch. + :return: A tuple of the per-batch results and whether the input was a single sentence. + """ + was_single = False + 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] + + sentence_batches = list(self._batch(sentences, batch_size)) + total_batches = math.ceil(len(sentences) / batch_size) + + self._set_max_length_in_tokenizer(max_length) + try: + 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(batch_fn)(batch, *batch_args) for batch in sentence_batches + ) + else: + results = [ + batch_fn(batch, *batch_args) + for batch in tqdm(sentence_batches, total=total_batches, disable=not show_progress_bar) + ] + finally: + self._set_max_length_in_tokenizer(self.max_length) + + return results, was_single + @overload def encode_as_sequence( self, @@ -295,41 +349,18 @@ def encode_as_sequence( :param multiprocessing_threshold: The threshold in number of sentences for using multiprocessing. :return: The encoded sentences with an embedding per token. """ - was_single = False - 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(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) + results, was_single = self._encode_dispatch( + sentences, + max_length=max_length, + batch_size=batch_size, + show_progress_bar=show_progress_bar, + use_multiprocessing=use_multiprocessing, + multiprocessing_threshold=multiprocessing_threshold, + batch_fn=self._encode_batch_as_sequence, + ) + out_array: list[np.ndarray] = [] + for r in results: + out_array.extend(r) if was_single: return out_array[0] @@ -380,45 +411,22 @@ def encode( :param **kwargs: Any additional arguments. These are ignored. :return: The encoded sentences. If a single sentence was passed, a vector is returned. """ - was_single = False - if isinstance(sentences, str): - sentences = [sentences] - was_single = True if isinstance(max_length, _UnsetType): 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) - 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) + results, was_single = self._encode_dispatch( + sentences, + max_length=max_length, + batch_size=batch_size, + show_progress_bar=show_progress_bar, + use_multiprocessing=use_multiprocessing, + multiprocessing_threshold=multiprocessing_threshold, + batch_fn=self._encode_batch, + batch_args=(normalize,), + ) + out_array = np.concatenate(results, axis=0) if was_single: return out_array[0]