huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 11:08:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: adds reset_kv_cache (#1335)

Browse Source
This commit is contained in:
drbh
2023-11-16 16:17:42 -05:00
committed by GitHub
parent 92a05b51cf
commit a1f41ab37b
1 changed files with 20 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
candle-transformers/src/models/trocr.rs
View File

@ -138,6 +138,10 @@ impl TrOCRAttention {
}) })
} }
fn reset_kv_cache(&mut self) {
self.kv_cache = None
}
fn _shape(&self, tensor: &Tensor, bsz: usize) -> Result<Tensor> { fn _shape(&self, tensor: &Tensor, bsz: usize) -> Result<Tensor> {
tensor tensor
.reshape((bsz, (), self.num_heads, self.head_dim))? .reshape((bsz, (), self.num_heads, self.head_dim))?
@ -239,6 +243,10 @@ impl TrOCRDecoderLayer {
}) })
} }
fn reset_kv_cache(&mut self) {
self.self_attn.reset_kv_cache();
}
fn forward( fn forward(
&mut self, &mut self,
xs: &Tensor, xs: &Tensor,
@ -307,6 +315,10 @@ impl TrOCRDecoder {
}) })
} }
fn reset_kv_cache(&mut self) {
self.layers.iter_mut().for_each(|l| l.reset_kv_cache())
}
pub fn forward( pub fn forward(
&mut self, &mut self,
xs: &Tensor, xs: &Tensor,
@ -393,6 +405,10 @@ impl TrOCRForCausalLM {
Ok(xs) Ok(xs)
} }
fn reset_kv_cache(&mut self) {
self.decoder.reset_kv_cache();
}
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@ -431,4 +447,8 @@ impl TrOCRModel {
self.decoder self.decoder
.forward(xs, Some(encoder_xs), past_kv_len, &mask) .forward(xs, Some(encoder_xs), past_kv_len, &mask)
} }
pub fn reset_kv_cache(&mut self) {
self.decoder.reset_kv_cache();
}
} }