Update preprocessing_molmo.py

preprocessing_molmo.py

@@ -97,7 +97,7 @@ class MolmoProcessor(ProcessorMixin):
             self._special_tokens = get_special_token_ids(self.tokenizer)
         return self._special_tokens
 
-    def get_tokens_input(self, prompt, message_format, always_start_with_space):
+    def get_tokens_input(self, prompt, message_format, always_start_with_space, out_text=None):
         if message_format == "none" or message_format is None:
             pass
         elif message_format == "role":
@@ -107,21 +107,9 @@ class MolmoProcessor(ProcessorMixin):
 
         if always_start_with_space:
             prompt = " " + prompt
-
-        tokens = self.tokenizer.encode(prompt, add_special_tokens=False)
-
-        return tokens
-
-    def get_tokens_input_for_logits(self, prompt, pred, message_format, always_start_with_space):
-        if message_format == "none" or message_format is None:
-            pass
-        elif message_format == "role":
-            prompt = "User: " + prompt + " Assistant: " + pred
-        else:
-            raise NotImplementedError(f"Message format {message_format} not implemented")
-
-        if always_start_with_space:
-            prompt = " " + prompt
+        
+        if out_text is not None:
+            prompt = " ".join([prompt, out_text])
 
         tokens = self.tokenizer.encode(prompt, add_special_tokens=False)
 
@@ -131,8 +119,10 @@ class MolmoProcessor(ProcessorMixin):
         self,
         text: TextInput = None,
         images: ImageInput = None,
+        out_text: TextInput = None,
         *,
         tokens: Optional[PreTokenizedInput] = None,
+        out_tokens: Optional[PreTokenizedInput] = None,
         **kwargs: Unpack[MolmoProcessorKwargs],
     ):
         output_kwargs = self._merge_kwargs(
@@ -146,8 +136,12 @@ class MolmoProcessor(ProcessorMixin):
                 text,
                 output_kwargs["text_kwargs"]["message_format"],
                 output_kwargs["text_kwargs"]["always_start_with_space"],
+                out_text
             )
 
+        if out_tokens is not None:
+            tokens = torch.cat([tokens, out_tokens], dim=0).tolist()
+
         image_token_id = self.special_token_ids[IMAGE_PROMPT]
 
         if images is not None: