fix: support get_current_tokenizer again

Author: percevalw

edsnlp/core/stream.py

@@ -66,6 +66,7 @@ def __bool__(self):
 
 
 INFER = _InferType()
+CONTEXT = [{}]
 
 T = TypeVar("T")
 
@@ -152,18 +153,23 @@ def __repr__(self):
 
 
 class MapOp(Op):
-    def __init__(self, pipe, kwargs):
+    def __init__(self, pipe, kwargs, context=None):
         self.pipe = pipe
         self.kwargs = kwargs
         self.is_generator = deep_isgeneratorfunction(pipe)
         self.elementwise = not self.is_generator
+        self.context = context or {}
 
     def __call__(self, items):
         for item in items:
             if isinstance(item, StreamSentinel):
                 yield item
                 continue
+
+            CONTEXT[0], old = self.context, CONTEXT[0]
             res = self.pipe(item, **self.kwargs)
+            CONTEXT[0] = old
+
             if self.is_generator:
                 yield from res
             else:
@@ -178,21 +184,24 @@ def __repr__(self):
 
 
 class MapBatchesOp(Op):
-    def __init__(self, pipe, kwargs, elementwise=False):
+    def __init__(self, pipe, kwargs, context=None, elementwise=False):
         self.pipe = pipe
         self.kwargs = kwargs
         self.is_generator = deep_isgeneratorfunction(pipe)
         if elementwise and self.is_generator:
             raise ValueError("Cannot use elementwise=True with a generator function")
         self.elementwise = elementwise
+        self.context = context or {}
 
     def __call__(self, batches):
         if hasattr(self.pipe, "batch_process"):
             for batch in batches:
                 if isinstance(batch, StreamSentinel):
                     yield batch
                     continue
+                CONTEXT[0], old = self.context, CONTEXT[0]
                 res = self.pipe.batch_process(batch, **self.kwargs)
+                CONTEXT[0] = old
                 res = list(res) if self.is_generator else (res,)
                 yield from res
         else:
@@ -202,11 +211,13 @@ def __call__(self, batches):
                     continue
                 results = []
                 for item in batch:
+                    CONTEXT[0], old = self.context, CONTEXT[0]
                     res = (
                         item
                         if isinstance(item, StreamSentinel)
                         else self.pipe(item, **self.kwargs)
                     )
+                    CONTEXT[0] = old
                     res = list(res) if self.is_generator else (res,)
                     results.extend(res)
                 yield results
@@ -727,6 +738,8 @@ def map_pipeline(
                 )
             ):
                 op.kwargs["tokenizer"] = tokenizer
+            if isinstance(op, (MapOp, MapBatchesOp)):
+                op.context["tokenizer"] = tokenizer
             new_ops.append(op)
         new_ops.append(MapOp(model._ensure_doc, {}))
         batch_size, batch_by = self.validate_batching(batch_size, batch_by)

edsnlp/data/converters.py

@@ -4,7 +4,6 @@
 Doc objects, and writers convert Doc objects to dictionaries.
 """
 
-import contextlib
 import inspect
 from copy import copy
 from types import FunctionType
@@ -26,6 +25,7 @@
 
 import edsnlp
 from edsnlp import registry
+from edsnlp.core.stream import CONTEXT
 from edsnlp.utils.bindings import BINDING_GETTERS
 from edsnlp.utils.span_getters import (
     SpanGetterArg,
@@ -136,21 +136,13 @@ def validate_attributes_mapping(value: AttributesMappingArg) -> Dict[str, str]:
 
 def get_current_tokenizer():
     global _DEFAULT_TOKENIZER
+    if "tokenizer" in CONTEXT[0]:
+        return CONTEXT[0]["tokenizer"]
     if _DEFAULT_TOKENIZER is None:
         _DEFAULT_TOKENIZER = edsnlp.blank("eds").tokenizer
     return _DEFAULT_TOKENIZER
 
 
-@contextlib.contextmanager
-def set_current_tokenizer(tokenizer):
-    global _DEFAULT_TOKENIZER
-    old = _DEFAULT_TOKENIZER
-    if tokenizer:
-        _DEFAULT_TOKENIZER = tokenizer
-    yield
-    _DEFAULT_TOKENIZER = old
-
-
 @registry.factory.register("eds.standoff_dict2doc", spacy_compatible=False)
 class StandoffDict2DocConverter:
     """