Get the attention mask for the given sequence length.
(&mut self, seq_len: usize, device: &Device)
| 148 | |
| 149 | /// Get the attention mask for the given sequence length. |
| 150 | pub fn mask(&mut self, seq_len: usize, device: &Device) -> Result<Tensor> { |
| 151 | // Always create/cache on self.device, then copy to target if needed |
| 152 | if !self.masks.contains_key(&seq_len) { |
| 153 | let mask: Vec<_> = (0..seq_len) |
| 154 | .flat_map(|i| (0..seq_len).map(move |j| u8::from(j > i))) |
| 155 | .collect(); |
| 156 | let mask = Tensor::from_slice(&mask, (seq_len, seq_len), &self.device)?; |
| 157 | self.masks.insert(seq_len, mask); |
| 158 | } |
| 159 | self.masks.get(&seq_len).unwrap().clone().to_device(device) |
| 160 | } |
| 161 | |
| 162 | /// Process the input k and v by either generating their cache entry or applying a previously cached one. |
| 163 | pub fn process_kv( |