amanrangapur syhuggingface commited on
Commit
ed442fa
·
verified ·
1 Parent(s): 0e72795

Update modeling_molmo.py for compatibility (#17)

Browse files

- Update modeling_molmo.py for compatibility (1f923ce84df12bd76bd65eb029df392543334159)


Co-authored-by: Seungyeon Jwa <syhuggingface@users.noreply.huggingface.co>

Files changed (1) hide show
  1. modeling_molmo.py +5 -0
modeling_molmo.py CHANGED
@@ -2273,6 +2273,11 @@ class MolmoForCausalLM(PreTrainedModel):
2273
  del model_kwargs["image_masks"]
2274
  del model_kwargs["image_input_idx"]
2275
  cache_name, cache = super()._extract_past_from_model_output(outputs)
 
 
 
 
 
2276
  model_kwargs[cache_name] = cache
2277
  model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
2278
  return model_kwargs
 
2273
  del model_kwargs["image_masks"]
2274
  del model_kwargs["image_input_idx"]
2275
  cache_name, cache = super()._extract_past_from_model_output(outputs)
2276
+ try:
2277
+ cache_name, cache = super()._extract_past_from_model_output(outputs)
2278
+ except AttributeError:
2279
+ past_key_values = outputs.past_key_values if "past_key_values" in outputs else None
2280
+ cache_name, cache = "past_key_values", past_key_values
2281
  model_kwargs[cache_name] = cache
2282
  model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
2283
  return model_kwargs