Updates the cache with the new `recurrent_state`/`attn_state`/`conv_state` for the layer `layer_idx`. Args: recurrent_state (`torch.Tensor`, `optional`): The new recurrent state to cache. attn_state (`Tuple[torch.Tensor, torch.Tensor]`, `opti
(
self,
recurrent_state: torch.Tensor = None,
attn_state: Tuple[torch.Tensor, torch.Tensor] = None,
conv_state: Tuple[torch.Tensor] = None,
layer_idx: int = 0,
offset: Optional[int] = 1,
cache_kwargs: Optional[Dict[str, Any]] = None,
)
| 38 | return len(self.states) |
| 39 | |
| 40 | def update( |
| 41 | self, |
| 42 | recurrent_state: torch.Tensor = None, |
| 43 | attn_state: Tuple[torch.Tensor, torch.Tensor] = None, |
| 44 | conv_state: Tuple[torch.Tensor] = None, |
| 45 | layer_idx: int = 0, |
| 46 | offset: Optional[int] = 1, |
| 47 | cache_kwargs: Optional[Dict[str, Any]] = None, |
| 48 | ) -> Dict[str, Any]: |
| 49 | """ |
| 50 | Updates the cache with the new `recurrent_state`/`attn_state`/`conv_state` for the layer `layer_idx`. |
| 51 | |
| 52 | Args: |
| 53 | recurrent_state (`torch.Tensor`, `optional`): |
| 54 | The new recurrent state to cache. |
| 55 | attn_state (`Tuple[torch.Tensor, torch.Tensor]`, `optional`): |
| 56 | The new attention key/value states to cache. |
| 57 | conv_state (`Tuple[torch.Tensor]`, `optional`): |
| 58 | The new convolution state to cache. |
| 59 | layer_idx (`int`, defaults to 0): |
| 60 | The index of the layer to cache the states for. |
| 61 | offset (`int`, `optional`, defaults to 1): |
| 62 | The number of new tokens being processed. |
| 63 | cache_kwargs (`Dict[str, Any]`, `optional`): |
| 64 | Additional arguments for the cache subclass. |
| 65 | |
| 66 | Return: |
| 67 | Dictionary of the updated state. |
| 68 | """ |
| 69 | # Update the number of seen tokens |
| 70 | if layer_idx == 0: |
| 71 | self._seen_tokens += offset |
| 72 | |
| 73 | if attn_state is not None: |
| 74 | # shape of `[b h l d]`` |
| 75 | input_size = attn_state[0].shape[-2] |
| 76 | window_size = cache_kwargs.get('window_size', None) |
| 77 | if not isinstance(attn_state, Tuple) or len(attn_state) != 2: |
| 78 | raise ValueError("`attn_state` must be a tuple of two tensors for key/value states") |
| 79 | if len(self.states) <= layer_idx: |
| 80 | if attn_state is not None: |
| 81 | if window_size is not None and input_size > window_size: |
| 82 | attn_state = (attn_state[0][..., -window_size:, :].contiguous(), |
| 83 | attn_state[1][..., -window_size:, :].contiguous()) |
| 84 | state = dict( |
| 85 | recurrent_state=recurrent_state, |
| 86 | attn_state=attn_state, |
| 87 | conv_state=conv_state, |
| 88 | ) |
| 89 | self.states.append(state) |
| 90 | else: |
| 91 | state = self.states[layer_idx] |
| 92 | if recurrent_state is not None: |
| 93 | state['recurrent_state'] = recurrent_state |
| 94 | if attn_state is not None: |
| 95 | key_state, value_state = state['attn_state'] |
| 96 | if window_size is not None and key_state.shape[-2] == window_size: |
| 97 | # DO NOT allocate new memory if the cache is full |
no test coverage detected