Fix thread-unsafe fast tokenizer usage in Tokenize

src/mgds/pipelineModules/Tokenize.py

@@ -1,3 +1,5 @@
+import threading
+
 import torch
 from transformers import CLIPTokenizer, T5Tokenizer, T5TokenizerFast, GemmaTokenizer, LlamaTokenizer, Qwen2Tokenizer, LlamaTokenizerFast
 
@@ -40,6 +42,15 @@ def __init__(
         self.suffix_text = suffix_text
         self.expand_mask = expand_mask
 
+        # fast tokenizers mutate shared Rust-side state (eg set_truncation_and_padding) on every
+        # call, which isn't safe under concurrent use of the same tokenizer instance from multiple
+        # dataloader threads and can raise "RuntimeError: Already borrowed". The lock is stored on
+        # the tokenizer itself so it's shared by every Tokenize instance wrapping that tokenizer.
+        # workaround for https://github.com/huggingface/transformers/issues/47085
+        if not hasattr(tokenizer, "_mgds_tokenizer_lock"):
+            tokenizer._mgds_tokenizer_lock = threading.Lock()
+        self.tokenizer_lock = tokenizer._mgds_tokenizer_lock
+
     def length(self) -> int:
         return self._get_previous_length(self.in_name)
 
@@ -58,31 +69,32 @@ def get_item(self, variation: int, index: int, requested_name: str = None) -> di
             text = self.format_text.format(text)
             max_length += self.additional_format_text_tokens
 
-        if self.apply_chat_template is not None:
-            messages = self.apply_chat_template(text)
-            text = self.tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                **self.apply_chat_template_kwargs,
+        with self.tokenizer_lock:
+            if self.apply_chat_template is not None:
+                messages = self.apply_chat_template(text)
+                text = self.tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    **self.apply_chat_template_kwargs,
+                )
+
+            tokenizer_output = self.tokenizer(
+                text,
+                padding='max_length',
+                truncation=True,
+                max_length=max_length,
+                return_tensors="pt",
             )
 
-        tokenizer_output = self.tokenizer(
-            text,
-            padding='max_length',
-            truncation=True,
-            max_length=max_length,
-            return_tensors="pt",
-        )
-
-        tokens = tokenizer_output.input_ids.to(self.pipeline.device)
-        mask = tokenizer_output.attention_mask.to(self.pipeline.device)
-
-        if self.suffix_text is not None:
-            suffix_output = self.tokenizer(self.suffix_text, return_tensors="pt")
-            suffix_tokens = suffix_output.input_ids.to(self.pipeline.device)
-            suffix_mask = suffix_output.attention_mask.to(self.pipeline.device)
-            tokens = torch.cat([tokens, suffix_tokens], dim=1)
-            mask = torch.cat([mask, suffix_mask], dim=1)
+            tokens = tokenizer_output.input_ids.to(self.pipeline.device)
+            mask = tokenizer_output.attention_mask.to(self.pipeline.device)
+
+            if self.suffix_text is not None:
+                suffix_output = self.tokenizer(self.suffix_text, return_tensors="pt")
+                suffix_tokens = suffix_output.input_ids.to(self.pipeline.device)
+                suffix_mask = suffix_output.attention_mask.to(self.pipeline.device)
+                tokens = torch.cat([tokens, suffix_tokens], dim=1)
+                mask = torch.cat([mask, suffix_mask], dim=1)
 
         tokens = tokens.squeeze(dim=0)
         mask = mask.squeeze(dim=0)