root commited on
Commit
69e3445
·
1 Parent(s): 03d86d0

Updated modeling_molmo.py for compatibility with transformers v>4.49

Browse files
Files changed (1) hide show
  1. modeling_molmoe.py +6 -1
modeling_molmoe.py CHANGED
@@ -2468,7 +2468,12 @@ class MolmoForCausalLM(PreTrainedModel):
2468
  del model_kwargs["images"]
2469
  del model_kwargs["image_masks"]
2470
  del model_kwargs["image_input_idx"]
2471
- cache_name, cache = super()._extract_past_from_model_output(outputs)
 
 
 
 
 
2472
  model_kwargs[cache_name] = cache
2473
  model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
2474
  return model_kwargs
 
2468
  del model_kwargs["images"]
2469
  del model_kwargs["image_masks"]
2470
  del model_kwargs["image_input_idx"]
2471
+ try:
2472
+ cache_name, cache = super()._extract_past_from_model_output(outputs)
2473
+ except AttributeError:
2474
+ past_key_values = outputs.past_key_values if "past_key_values" in outputs else None
2475
+ cache_name, cache = "past_key_values", past_key_values
2476
+
2477
  model_kwargs[cache_name] = cache
2478
  model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
2479
  return model_kwargs