Update preprocessing_molmo.py

preprocessing_molmo.py

@@ -97,22 +97,31 @@ 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, out_text=None):
+    def get_tokens_input(self, prompt, message_format, always_start_with_space, out_text=None, pad_length=None):
         if message_format == "none" or message_format is None:
             pass
         elif message_format == "role":
             prompt = "User: " + prompt + " Assistant:"
         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)
 
+        if pad_length is not None:
+            assert isinstance(pad_length, int)
+            if len(tokens) > pad_length:
+                tokens = tokens[:pad_length]
+        
+            if len(tokens) < pad_length:
+                pad_token_id = self.tokenizer.pad_token_id or 0  # Use 0 if pad_token_id is not set
+                tokens = tokens + [pad_token_id] * (pad_length - len(tokens))
+    
         return tokens
 
     def process(
@@ -120,6 +129,7 @@ class MolmoProcessor(ProcessorMixin):
         text: TextInput = None,
         images: ImageInput = None,
         out_text: TextInput = None,
+        pad_length: int = None, 
         *,
         tokens: Optional[PreTokenizedInput] = None,
         out_tokens: Optional[PreTokenizedInput] = None,
@@ -136,7 +146,8 @@ class MolmoProcessor(ProcessorMixin):
                 text,
                 output_kwargs["text_kwargs"]["message_format"],
                 output_kwargs["text_kwargs"]["always_start_with_space"],
-                out_text
+                out_text,
+                pad_length,
             )
 
         if out_tokens is not None: