Update tokenize.py for common >= 1.4.0

mistralai · Sep 13, 2024 · 8520eb8 · 8520eb8
1 parent 41c793f
commit 8520eb8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/finetune/data/tokenize.py b/finetune/data/tokenize.py
@@ -311,6 +311,10 @@ def tokenize_instruct(
                 is_first=msg_idx == first_user_idx,
                 system_prompt=sample.system_prompt,
             )
+            if isinstance(curr_tokens, tuple):
+                # Versions of mistral_common>1.3.4 return a tuple of tokens (text), tokens (image), spans (image)
+                curr_tokens = curr_tokens[0]
+
             curr_masks = [False] * len(curr_tokens)  # only predict bot answers
         elif isinstance(message, ToolMessage):
             curr_tokens = instruct_tokenizer.encode_tool_message(