huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add the sliding window. (#986)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-28 18:26:33 +02:00
committed by GitHub
parent 716ab2ccdc
commit 23b3576c47
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
candle-transformers/src/models/mistral.rs
View File

@ -299,7 +299,6 @@ pub struct Model {
layers: Vec<DecoderLayer>, layers: Vec<DecoderLayer>,
norm: RmsNorm, norm: RmsNorm,
lm_head: Linear, lm_head: Linear,
#[allow(unused)]
sliding_window: usize, sliding_window: usize,
device: Device, device: Device,
dtype: DType, dtype: DType,
@ -338,7 +337,15 @@ impl Model {
) -> Result<Tensor> { ) -> Result<Tensor> {
// Sliding window mask? // Sliding window mask?
let mask: Vec<_> = (0..tgt_len) let mask: Vec<_> = (0..tgt_len)
.flat_map(|i| (0..tgt_len).map(move |j| if j > i { f32::NEG_INFINITY } else { 0. })) .flat_map(|i| {
(0..tgt_len).map(move |j| {
if i < j || j + self.sliding_window < i {
f32::NEG_INFINITY
} else {
0.
}
})
})
.collect(); .collect();
let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?; let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), &self.device)?;
let mask = if seqlen_offset > 0 { let mask = if seqlen_offset > 0 {