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 some repeat penalty to the phi example. (#961)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-25 20:53:30 +01:00
committed by GitHub
parent a36d883254
commit c78a294323
1 changed files with 40 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
candle-examples/examples/phi/main.rs
View File

@ -26,15 +26,20 @@ struct TextGeneration {
device: Device, device: Device,
tokenizer: Tokenizer, tokenizer: Tokenizer,
logits_processor: LogitsProcessor, logits_processor: LogitsProcessor,
repeat_penalty: f32,
repeat_last_n: usize,
} }
impl TextGeneration { impl TextGeneration {
#[allow(clippy::too_many_arguments)]
fn new( fn new(
model: Model, model: Model,
tokenizer: Tokenizer, tokenizer: Tokenizer,
seed: u64, seed: u64,
temp: Option<f64>, temp: Option<f64>,
top_p: Option<f64>, top_p: Option<f64>,
repeat_penalty: f32,
repeat_last_n: usize,
device: &Device, device: &Device,
) -> Self { ) -> Self {
let logits_processor = LogitsProcessor::new(seed, temp, top_p); let logits_processor = LogitsProcessor::new(seed, temp, top_p);
@ -42,6 +47,8 @@ impl TextGeneration {
model, model,
tokenizer, tokenizer,
logits_processor, logits_processor,
repeat_penalty,
repeat_last_n,
device: device.clone(), device: device.clone(),
} }
} }
@ -69,6 +76,16 @@ impl TextGeneration {
Model::Quantized(m) => m.forward(&input)?, Model::Quantized(m) => m.forward(&input)?,
}; };
let logits = logits.squeeze(0)?.to_dtype(DType::F32)?; let logits = logits.squeeze(0)?.to_dtype(DType::F32)?;
let logits = if self.repeat_penalty == 1. {
logits
} else {
let start_at = tokens.len().saturating_sub(self.repeat_last_n);
candle_transformers::utils::apply_repeat_penalty(
&logits,
self.repeat_penalty,
&tokens[start_at..],
)?
};
let next_token = self.logits_processor.sample(&logits)?; let next_token = self.logits_processor.sample(&logits)?;
tokens.push(next_token); tokens.push(next_token);
@ -127,6 +144,14 @@ struct Args {
#[arg(long)] #[arg(long)]
quantized: bool, quantized: bool,
/// Penalty to be applied for repeating tokens, 1. means no penalty.
#[arg(long, default_value_t = 1.1)]
repeat_penalty: f32,
/// The context size to consider for the repeat penalty.
#[arg(long, default_value_t = 64)]
repeat_last_n: usize,
} }
fn main() -> Result<()> { fn main() -> Result<()> {
@ -134,7 +159,6 @@ fn main() -> Result<()> {
use tracing_subscriber::prelude::*; use tracing_subscriber::prelude::*;
let args = Args::parse(); let args = Args::parse();
let _guard = if args.tracing { let _guard = if args.tracing {
let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
tracing_subscriber::registry().with(chrome_layer).init(); tracing_subscriber::registry().with(chrome_layer).init();
@ -142,6 +166,19 @@ fn main() -> Result<()> {
} else { } else {
None None
}; };
println!(
"avx: {}, neon: {}, simd128: {}, f16c: {}",
candle::utils::with_avx(),
candle::utils::with_neon(),
candle::utils::with_simd128(),
candle::utils::with_f16c()
);
println!(
"temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}",
args.temperature.unwrap_or(0.),
args.repeat_penalty,
args.repeat_last_n
);
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let api = Api::new()?; let api = Api::new()?;
@ -185,6 +222,8 @@ fn main() -> Result<()> {
args.seed, args.seed,
args.temperature, args.temperature,
args.top_p, args.top_p,
args.repeat_penalty,
args.repeat_last_n,
&device, &device,
); );
pipeline.run(&args.prompt, args.sample_len)?; pipeline.run(&args.prompt, args.sample_len)?;