Fix cache type in Idefics2 (#30729)

standardize cache in idefics2
This commit is contained in:
Raushan Turganbay 2024-05-14 13:30:53 +05:00 committed by GitHub
parent 449894d2e5
commit c02d302e6b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 3 deletions

View File

@ -1591,9 +1591,10 @@ class Idefics2Model(Idefics2PreTrainedModel):
raise ValueError("You have to specify either input_ids or inputs_embeds")
past_seen_tokens = 0
if use_cache:
if not isinstance(past_key_values, Cache):
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
return_legacy_cache = False
if use_cache and not isinstance(past_key_values, Cache): # kept for BC (non `Cache` `past_key_values` inputs)
return_legacy_cache = True
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
past_seen_tokens = past_key_values.get_usable_length(seq_length)
if inputs_embeds is not None and input_ids is None and past_seen_tokens == 0:
@ -1667,6 +1668,9 @@ class Idefics2Model(Idefics2PreTrainedModel):
return_dict=return_dict,
)
if return_legacy_cache:
outputs.past_key_values = outputs.past_key_values.to_legacy_cache()
if not return_dict:
return tuple(v for v in [*outputs, image_hidden_states] if v is not None)