oumi-ai
/

Molmo-7B-O-0924

Image-Text-to-Text

text-generation

Model card Files Files and versions

Oussama Elachqar commited on Jul 8

Commit

37ff428

·

1 Parent(s): f19f7fc

fix cache

Files changed (1) hide show

modeling_molmo.py +5 -1

modeling_molmo.py CHANGED Viewed

@@ -2275,7 +2275,11 @@ class MolmoForCausalLM(PreTrainedModel):
                 del model_kwargs["images"]
                 del model_kwargs["image_masks"]
                 del model_kwargs["image_input_idx"]
-        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

                 del model_kwargs["images"]
                 del model_kwargs["image_masks"]
                 del model_kwargs["image_input_idx"]
+        try:
+            cache_name, cache = super()._extract_past_from_model_output(outputs)
+        except AttributeError:
+            past_key_values = outputs.past_key_values if "past_key_values" in outputs else None
+            cache_name, cache = "past_key_values", past_key_values
         model_kwargs[cache_name] = cache
         model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
         return model_kwargs