Skip to content

Fix past_key_values type hint in model output types #37953

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gante merged 5 commits into from
May 13, 2025
Merged

Fix past_key_values type hint in model output types #37953

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
268eb5d
F: Fix type hint.
ChengLyu May 5, 2025
c06d062
F: Use Cache type.
ChengLyu May 5, 2025
4386271
F: Sort import.
ChengLyu May 5, 2025
92e7c2c
U: Format.
ChengLyu May 5, 2025
8bb4f48
U: Address reviews.
ChengLyu May 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/transformers/cache_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -464,7 +464,7 @@ def get_max_cache_shape(self) -> Optional[int]:
"""Returns the maximum sequence length of the cache object. DynamicCache does not have a maximum length."""
return None

def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]:
def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
"""Converts the `DynamicCache` instance into the its equivalent in the legacy cache format. Used for
backward compatibility."""
legacy_cache = ()
Expand All @@ -473,7 +473,9 @@ def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]:
return legacy_cache

@classmethod
def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "DynamicCache":
def from_legacy_cache(
cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor, torch.FloatTensor]]] = None
) -> "DynamicCache":
"""Converts a cache in the legacy cache format into an equivalent `DynamicCache`. Used for
backward compatibility."""
cache = cls()
Expand Down Expand Up @@ -1505,8 +1507,8 @@ def __len__(self):
"""
return len(self.self_attention_cache)

def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]:
"""Converts the `EncoderDecoderCache` instance into its equivalent in the legacy cache format."""
def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor]]:
"""Converts the `EncoderDecoderCache` instance into its equivalent in the legacy cache format."""
legacy_cache = ()
if len(self.cross_attention_cache) > 0:
for self_attn, cross_attn in zip(
Expand Down
Loading