From 7f0ad6c2ce599d5fe71d6c678cc2b829ffaa98ef Mon Sep 17 00:00:00 2001
From: stephantul <stephantul@gmail.com>
Date: Thu, 5 Nov 2026 16:24:28 +0100
Subject: [PATCH] fix: parallelism environment variable should be reset

---
 model2vec/model.py  |  5 +++++
 tests/test_model.py | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+)

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: