Update modeling_molmo.py

modeling_molmo.py

@@ -2509,7 +2509,6 @@ class MOLMoForCausalLM(PreTrainedModel):
         outputs: ModelOutput,
         model_kwargs: Dict[str, Any],
         is_encoder_decoder: bool = False,
-        standardize_cache_format: bool = False,
         num_new_tokens: int = 1,
     ) -> Dict[str, Any]:
         if self.config.use_position_ids:
@@ -2520,8 +2519,9 @@ class MOLMoForCausalLM(PreTrainedModel):
                 del model_kwargs["images"]
                 del model_kwargs["image_masks"]
                 del model_kwargs["image_input_idx"]
-        model_kwargs = super()._update_model_kwargs_for_generation(
-            outputs, model_kwargs, is_encoder_decoder, standardize_cache_format, num_new_tokens)
+        cache_name, cache = super()._extract_past_from_model_output(outputs)
+        model_kwargs[cache_name] = cache
+        model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
         return model_kwargs
 
     # TODO: these are required to make the implementation complete.