diff --git a/model2vec/model.py b/model2vec/model.py index 2a3f6c5..6f43098 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -276,6 +276,7 @@ def _encode_dispatch( total_batches = math.ceil(len(sentences) / batch_size) self._set_max_length_in_tokenizer(max_length) + previous_parallelism = os.environ.get("TOKENIZERS_PARALLELISM") try: if use_multiprocessing and len(sentences) > multiprocessing_threshold: # Disable parallelism for tokenizers @@ -291,6 +292,10 @@ def _encode_dispatch( ] finally: self._set_max_length_in_tokenizer(self.max_length) + if previous_parallelism is None: + os.environ.pop("TOKENIZERS_PARALLELISM", None) + else: + os.environ["TOKENIZERS_PARALLELISM"] = previous_parallelism return results, was_single diff --git a/tests/test_model.py b/tests/test_model.py index b7a062a..aaf5866 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -1,3 +1,4 @@ +import os from pathlib import Path import numpy as np @@ -112,6 +113,24 @@ def test_encode_multiprocessing( assert encoded.shape == (15000, 2) +@pytest.mark.parametrize("previous", [None, "true"]) +def test_encode_multiprocessing_restores_tokenizers_parallelism( + mock_vectors: np.ndarray, + mock_tokenizer: Tokenizer, + mock_config: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + previous: str | None, +) -> None: + """Test that multiprocessing encoding restores TOKENIZERS_PARALLELISM.""" + if previous is None: + monkeypatch.delenv("TOKENIZERS_PARALLELISM", raising=False) + else: + monkeypatch.setenv("TOKENIZERS_PARALLELISM", previous) + model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config) + model.encode(["word1 word2"] * 15_000, use_multiprocessing=True) + assert os.environ.get("TOKENIZERS_PARALLELISM") == previous + + def test_encode_as_sequence_max_length( mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str] ) -> None: