Update preprocessing_molmo.py

preprocessing_molmo.py

@@ -23,7 +23,7 @@ from transformers.processing_utils import (
     ProcessorMixin,
 )
 
-from transformers.tokenization_utils_base import TextInput
+from transformers.tokenization_utils_base import TextInput, PreTokenizedInput
 from transformers.utils import logging
 
 from transformers import AutoTokenizer
@@ -116,6 +116,8 @@ class MolmoProcessor(ProcessorMixin):
         self,
         text: TextInput = None,
         images: ImageInput = None,
+        *,
+        tokens: Optional[PreTokenizedInput] = None,
         **kwargs: Unpack[MolmoProcessorKwargs],
     ):
         output_kwargs = self._merge_kwargs(
@@ -124,11 +126,12 @@ class MolmoProcessor(ProcessorMixin):
             **kwargs,
         )
 
-        tokens = self.get_tokens_input(
-            text,
-            output_kwargs["text_kwargs"]["message_format"],
-            output_kwargs["text_kwargs"]["always_start_with_space"],
-        )
+        if tokens is None:
+            tokens = self.get_tokens_input(
+                text,
+                output_kwargs["text_kwargs"]["message_format"],
+                output_kwargs["text_kwargs"]["always_start_with_space"],
+            )
 
         image_token_id = self.special_token_ids[IMAGE_PROMPT]