fix: parallelism environment variable should be reset

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
 

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: