File size: 3,584 Bytes
fe0450e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
# -*- coding: utf-8 -*-
from __future__ import annotations
from typing import Any, Dict, List, Optional, Tuple
import torch
import transformers
class Cache(transformers.cache_utils.Cache):
"""
A cache used for storing hidden states produced by flash linear attention models.
It stores the states of each layer as the tensor of shape `[batch_size, key_dim, value_dim]`.
"""
def __init__(
self,
seen_tokens: int = 0
) -> Cache:
self.states: List[torch.Tensor] = []
self._seen_tokens = seen_tokens # Used in `generate` to keep tally of how many tokens the cache has seen
def __getitem__(self, layer_idx: int) -> torch.Tensor:
if layer_idx < len(self):
return self.states[layer_idx]
else:
raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")
def __iter__(self):
for state in self.states:
yield state
def __len__(self):
return len(self.states)
def update(
self,
state: Tuple[torch.Tensor],
layer_idx: int,
offset: Optional[int] = 1,
cache_kwargs: Optional[Dict[str, Any]] = None,
) -> Tuple[torch.Tensor]:
"""
Updates the cache with the new `state` for the layer `layer_idx`.
Parameters:
state (`Tuple[torch.Tensor]`):
The new state to cache.
layer_idx (`int`):
The index of the layer to cache the states for.
offset (`int`):
The offset of current fed tokens.
cache_kwargs (`Dict[str, Any]`, `optional`):
Additional arguments for the cache subclass.
Return:
The updated state.
"""
if isinstance(state, torch.Tensor):
state = (state,)
if len(self.states) <= layer_idx:
self.states.append(state)
else:
for i, s in enumerate(state):
self.states[layer_idx][i].copy_(s)
# update the number of seen tokens once we achieve the last layer
if layer_idx == len(self) - 1:
self._seen_tokens += offset
return state
def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
"""Returns the sequence length of the cached states. A layer index can be optionally passed."""
if len(self.states) <= layer_idx:
return 0
return self._seen_tokens
def get_max_length(self) -> Optional[int]:
"""Returns the maximum sequence length of the cached states. Cache does not have a maximum length."""
return None
def reorder_cache(self, beam_idx: torch.LongTensor):
"""Reorders the cache for beam search, given the selected beam indices."""
for layer_idx in range(len(self.states)):
device = self.states[layer_idx].device
self.states[layer_idx] = self.states[layer_idx].index_select(0, beam_idx.to(device))
def to_legacy_cache(self) -> Tuple[torch.Tensor]:
return tuple(self.states)
@classmethod
def from_legacy_cache(
cls,
past_key_values: Optional[Tuple[torch.Tensor]] = None,
seen_tokens: int = 0
) -> Cache:
"""Converts a cache in the legacy cache format into an equivalent `Cache`."""
cache = cls(seen_tokens)
if past_key_values is not None:
for layer_idx in range(len(past_key_values)):
cache.update(past_key_values[layer_idx], layer_idx)
return cache
|