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

Allow for arbitrary temperature modifications.

Browse Source
This commit is contained in:
laurent
2024-03-23 15:47:39 +01:00
parent a62a97340c
commit 5e70821dd0
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-transformers/src/generation/mod.rs
View File

@ -87,11 +87,17 @@ impl LogitsProcessor {
} }
pub fn sample(&mut self, logits: &Tensor) -> Result<u32> { pub fn sample(&mut self, logits: &Tensor) -> Result<u32> {
self.sample_f(logits, |_| {})
}
pub fn sample_f(&mut self, logits: &Tensor, f: impl FnOnce(&mut [f32])) -> Result<u32> {
let logits = logits.to_dtype(DType::F32)?; let logits = logits.to_dtype(DType::F32)?;
let prs = |temperature: f64| -> Result<Vec<f32>> { let prs = |temperature: f64| -> Result<Vec<f32>> {
let logits = (&logits / temperature)?; let logits = (&logits / temperature)?;
let prs = candle_nn::ops::softmax_last_dim(&logits)?; let prs = candle_nn::ops::softmax_last_dim(&logits)?;
prs.to_vec1() let mut prs = prs.to_vec1()?;
f(&mut prs);
Ok(prs)
}; };
let next_token = match &self.sampling { let next_token = match &self.sampling {