diff --git a/src/mgds/pipelineModules/Tokenize.py b/src/mgds/pipelineModules/Tokenize.py index 77a5a12..e0352b0 100644 --- a/src/mgds/pipelineModules/Tokenize.py +++ b/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)