huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 10:26:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Mandatory topk sampling for recurrent-gemma. (#2051)

Browse Source
This commit is contained in:
Laurent Mazare
2024-04-13 10:31:39 +02:00
committed by GitHub
parent 2bf413caa3
commit 26cbbf8d84
1 changed files with 20 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
candle-examples/examples/recurrent-gemma/main.rs
View File

@ -41,11 +41,26 @@ impl TextGeneration {
seed: u64,
temp: Option<f64>,
top_p: Option<f64>,
top_k: usize,
repeat_penalty: f32,
repeat_last_n: usize,
device: &Device,
) -> Self {
let logits_processor = LogitsProcessor::new(seed, temp, top_p);
let sampling = match temp {
None => candle_transformers::generation::Sampling::ArgMax,
Some(temperature) => match top_p {
None => candle_transformers::generation::Sampling::TopK {
temperature,
k: top_k,
},
Some(top_p) => candle_transformers::generation::Sampling::TopKThenTopP {
temperature,
k: top_k,
p: top_p,
},
},
};
let logits_processor = LogitsProcessor::from_sampling(seed, sampling);
Self {
model,
tokenizer: TokenOutputStream::new(tokenizer),
@ -143,6 +158,9 @@ struct Args {
#[arg(long)]
top_p: Option<f64>,
#[arg(long, default_value_t = 250)]
top_k: usize,
/// The seed to use when generating random samples.
#[arg(long, default_value_t = 299792458)]
seed: u64,
@ -256,6 +274,7 @@ fn main() -> Result<()> {
args.seed,
args.temperature,
args.top_p,
args.top_k,
args.repeat_penalty,
args.repeat_last_n,
&device,