tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

generate add penalty repeat

Browse Source
This commit is contained in:
jhqxxx
2026-04-03 14:15:10 +08:00
parent bd8eee6520
commit 279480e3d7
26 changed files with 132 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cli/mod.rs
+3
View File
@@ -240,6 +240,9 @@ pub(crate) fn run_run(args: RunArgs) -> anyhow::Result<()> {
WhichModel::Qwen3_0_6B => {
qwen3::Qwen3Exec::run(&input, output.as_deref(), &weight_path)?;
}
WhichModel::Qwen3_1_7B => {
qwen3::Qwen3Exec::run(&input, output.as_deref(), &weight_path)?;
}
WhichModel::Qwen3_5_0_8B => {
qwen3_5::Qwen3_5Exec::run(&input, output.as_deref(), &weight_path)?;
}
src/models/common/generate.rs
+27 -22
View File
@@ -49,6 +49,8 @@ pub fn get_logit_processor(
pub struct GenerationContext {
pub logit_processor: LogitsProcessor,
pub repeat_penalty: f32,
pub repeat_last_n: usize,
pub seqlen_offset: usize,
pub seq_len: usize,
pub sample_len: u32,
@@ -60,6 +62,8 @@ impl GenerationContext {
temperature: Option<f32>,
top_p: Option<f32>,
top_k: Option<usize>,
repeat_penalty: Option<f32>,
repeat_last_n: Option<usize>,
seed: u64,
initial_seq_len: usize,
max_tokens: u32,
@@ -67,6 +71,8 @@ impl GenerationContext {
) -> Self {
Self {
logit_processor: get_logit_processor(temperature, top_p, top_k, seed),
repeat_penalty: repeat_penalty.unwrap_or(1.0),
repeat_last_n: repeat_last_n.unwrap_or(64),
seqlen_offset: 0,
seq_len: initial_seq_len,
sample_len: max_tokens,
@@ -91,27 +97,27 @@ impl GenerationContext {
/// 采样辅助函数
fn sample_and_push(
processor: &mut LogitsProcessor,
ctx: &mut GenerationContext,
logits: &Tensor,
generated: &mut Vec<u32>,
) -> Result<u32> {
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
let token = processor.sample(&logits)?;
// 重复惩罚
let logits = if ctx.repeat_penalty == 1. || ctx.repeat_last_n == 0 {
logits
} else {
let start_at = generated.len().saturating_sub(ctx.repeat_last_n);
candle_transformers::utils::apply_repeat_penalty(
&logits,
ctx.repeat_penalty,
&generated[start_at..],
)?
};
let token = ctx.logit_processor.sample(&logits)?;
generated.push(token);
Ok(token)
}
// TODO
// let logits = if self.repeat_penalty == 1. {
// logits
// } else {
// let start_at = generate.len().saturating_sub(self.repeat_last_n);
// candle_transformers::utils::apply_repeat_penalty(
// &logits,
// self.repeat_penalty,
// &generate[start_at..],
// )?
// };
pub fn generate_generic<M: InferenceModel>(
model: &mut M,
tokenizer: &TokenizerModel,
@@ -125,7 +131,7 @@ pub fn generate_generic<M: InferenceModel>(
let eos_ids = model.stop_token_ids();
let i_start = Instant::now();
let logits = model.forward_initial(&input_ids, ctx.seqlen_offset, data)?;
let next_token = sample_and_push(&mut ctx.logit_processor, &logits, &mut generated)?;
let next_token = sample_and_push(ctx, &logits, &mut generated)?;
let i_duration = i_start.elapsed();
let prompt_secs = i_duration.as_secs_f64();
let mut input_ids = ctx.prepare_for_next_token(next_token)?;
@@ -134,12 +140,11 @@ pub fn generate_generic<M: InferenceModel>(
let i_start = Instant::now();
for _ in 1..ctx.sample_len {
let logits = model.forward_step(&input_ids, ctx.seqlen_offset)?;
let next_token = sample_and_push(&mut ctx.logit_processor, &logits, &mut generated)?;
let next_token = sample_and_push(ctx, &logits, &mut generated)?;
if eos_ids.contains(&next_token) {
break;
}
input_ids = ctx.prepare_for_next_token(next_token)?;
}
let i_duration = i_start.elapsed();
@@ -167,6 +172,8 @@ pub fn generate_stream_generic<M: InferenceModel>(
temperature: Option<f32>,
top_p: Option<f32>,
top_k: Option<usize>,
repeat_penalty: Option<f32>,
repeat_last_n: Option<usize>,
seed: u64,
max_tokens: u32,
in_reasoning: bool,
@@ -177,6 +184,8 @@ pub fn generate_stream_generic<M: InferenceModel>(
temperature,
top_p,
top_k,
repeat_penalty,
repeat_last_n,
seed,
input_ids.dim(1)?,
max_tokens,
@@ -193,7 +202,7 @@ pub fn generate_stream_generic<M: InferenceModel>(
let mut tool_call_id = None;
let mut tool_call_content = String::new();
let mut in_reasoning = in_reasoning;
// 处理 unicode 错误累积
let mut generated = Vec::new();
for _ in 0..ctx.sample_len {
let i_start = Instant::now();
let logits = if ctx.seqlen_offset == 0 {
@@ -202,11 +211,7 @@ pub fn generate_stream_generic<M: InferenceModel>(
} else {
model.forward_step(&input_ids, ctx.seqlen_offset)
}?;
let next_token = {
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
ctx.logit_processor.sample(&logits)?
};
let next_token = sample_and_push(&mut ctx, &logits, &mut generated)?;
completion_tokens += 1;
let i_duration = i_start.elapsed();
if ctx.seqlen_offset == 0 {
src/models/common/model_mapping.rs
+3
View File
@@ -18,6 +18,8 @@ pub enum WhichModel {
Qwen2_5VL7B,
#[value(name = "Qwen/Qwen3-0.6B")]
Qwen3_0_6B,
#[value(name = "Qwen/Qwen3-1.7B")]
Qwen3_1_7B,
#[value(name = "Qwen/Qwen3.5-0.8B")]
Qwen3_5_0_8B,
#[value(name = "Qwen/Qwen3.5-2B")]
@@ -93,6 +95,7 @@ impl WhichModel {
// LLM models
WhichModel::MiniCPM4_0_5B
| WhichModel::Qwen3_0_6B
| WhichModel::Qwen3_1_7B
| WhichModel::LFM2_1_2B
| WhichModel::LFM2_5_1_2BInstruct => "llm",
// VLM models
src/models/deepseek_ocr/generate.rs
+4
View File
@@ -98,6 +98,8 @@ impl GenerateModel for DeepseekOCRGenerateModel {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
mes.seed.unwrap_or(34562) as u64,
input_ids.dim(1)?,
max_tokens,
@@ -169,6 +171,8 @@ impl GenerateModel for DeepseekOCRGenerateModel {
temperature,
top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
max_tokens,
false,
src/models/fun_asr_nano/generate.rs
+4 -2
View File
@@ -109,6 +109,8 @@ impl GenerateModel for FunAsrNanoGenerateModel {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
max_tokens,
@@ -145,8 +147,6 @@ impl GenerateModel for FunAsrNanoGenerateModel {
let top_k = self.generation_config.top_k;
let seed = mes.seed.unwrap_or(34562) as u64;
let max_tokens = mes.max_tokens.unwrap_or(1024);
// let mut logit_processor =
// get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed);
let (speech, fbank_mask, input_ids) = self.processor.process_info(&mes, &self.tokenizer)?;
let speech = speech.to_dtype(self.dtype)?;
let data_vec = vec![speech.into(), fbank_mask.into()];
@@ -159,6 +159,8 @@ impl GenerateModel for FunAsrNanoGenerateModel {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
max_tokens,
false,
src/models/glm_asr_nano/generate.rs
+4
View File
@@ -77,6 +77,8 @@ impl<'a> GenerateModel for GlmAsrNanoGenerateModel<'a> {
mes.temperature,
mes.top_p,
mes.top_k,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -124,6 +126,8 @@ impl<'a> GenerateModel for GlmAsrNanoGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/glm_ocr/generate.rs
+4
View File
@@ -113,6 +113,8 @@ impl GenerateModel for GlmOcrGenerateModel {
mes.temperature,
mes.top_p,
mes.top_k,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -187,6 +189,8 @@ impl GenerateModel for GlmOcrGenerateModel {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/hunyuan_ocr/generate.rs
+4
View File
@@ -87,6 +87,8 @@ impl<'a> GenerateModel for HunyuanOCRGenerateModel<'a> {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -149,6 +151,8 @@ impl<'a> GenerateModel for HunyuanOCRGenerateModel<'a> {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/lfm2/generate.rs
+4
View File
@@ -72,6 +72,8 @@ impl<'a> GenerateModel for Lfm2GenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -114,6 +116,8 @@ impl<'a> GenerateModel for Lfm2GenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/lfm2vl/generate.rs
+4
View File
@@ -72,6 +72,8 @@ impl<'a> GenerateModel for Lfm2VLGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -126,6 +128,8 @@ impl<'a> GenerateModel for Lfm2VLGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/minicpm4/generate.rs
+4
View File
@@ -64,6 +64,8 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -104,6 +106,8 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/mod.rs
+4
View File
@@ -168,6 +168,10 @@ pub fn load_model<'a>(
let model = Qwen3GenerateModel::init(path, None, None)?;
ModelInstance::Qwen3(model)
}
WhichModel::Qwen3_1_7B => {
let model = Qwen3GenerateModel::init(path, None, None)?;
ModelInstance::Qwen3(model)
}
WhichModel::Qwen3_5_0_8B => {
let model = Qwen3_5GenerateModel::init(path, None, None)?;
ModelInstance::Qwen3_5(model)
src/models/paddleocr_vl/generate.rs
+4
View File
@@ -79,6 +79,8 @@ impl<'a> GenerateModel for PaddleOCRVLGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -141,6 +143,8 @@ impl<'a> GenerateModel for PaddleOCRVLGenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
false,
src/models/qwen3/generate.rs
+4
View File
@@ -71,6 +71,8 @@ impl<'a> GenerateModel for Qwen3GenerateModel<'a> {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -117,6 +119,8 @@ impl<'a> GenerateModel for Qwen3GenerateModel<'a> {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
in_reasoning,
src/models/qwen3_5/generate.rs
+10 -6
View File
@@ -29,8 +29,8 @@ pub struct Qwen3_5GenerateModel<'a> {
qwen3_5: Qwen3_5Model,
device: Device,
model_name: String,
// repeat_penalty: f32, // TODO
// repeat_last_n: usize,
repeat_penalty: f32,
repeat_last_n: usize,
}
impl<'a> Qwen3_5GenerateModel<'a> {
@@ -59,8 +59,8 @@ impl<'a> Qwen3_5GenerateModel<'a> {
qwen3_5,
device,
model_name: model_name.to_string(),
// repeat_penalty: 1.01,
// repeat_last_n: 64,
repeat_penalty: 1.01,
repeat_last_n: 64,
})
}
@@ -116,8 +116,8 @@ impl<'a> Qwen3_5GenerateModel<'a> {
qwen3_5,
device,
model_name: stem.to_string(),
// repeat_penalty: 1.1,
// repeat_last_n: 64,
repeat_penalty: 1.1,
repeat_last_n: 64,
})
}
}
@@ -147,6 +147,8 @@ impl<'a> GenerateModel for Qwen3_5GenerateModel<'a> {
temperature.into(),
top_p.into(),
Some(20),
self.repeat_penalty.into(),
self.repeat_last_n.into(),
seed,
input_ids.dim(1)?,
sample_len,
@@ -213,6 +215,8 @@ impl<'a> GenerateModel for Qwen3_5GenerateModel<'a> {
mes.temperature,
mes.top_p,
None,
self.repeat_penalty.into(),
self.repeat_last_n.into(),
seed,
sample_len,
in_reasoning,
src/models/qwen3vl/generate.rs
+4 -103
View File
@@ -84,6 +84,8 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
input_ids.dim(1)?,
sample_len,
@@ -150,6 +152,8 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
temperature.into(),
top_p.into(),
top_k.into(),
mes.repeat_penalty,
mes.repeat_last_n,
seed,
sample_len,
in_reasoning,
@@ -157,108 +161,5 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
&self.model_name,
)?;
Ok(Box::new(Box::pin(stream)))
// let sample_len = mes.max_tokens.unwrap_or(1024);
// let stream = stream! {
// let mut error_tokens = Vec::new();
// let mut pixel_values = input.pixel_values.as_ref();
// let image_grid_thw = input.image_grid_thw.as_ref();
// let mut pixel_values_video = input.pixel_values_video.as_ref();
// let video_grid_thw = input.video_grid_thw.as_ref();
// let mut tool_call_id = None;
// let mut tool_call_content = String::new();
// for _ in 0..sample_len {
// let logits = self.qwen3_vl.forward(
// &input_ids,
// pixel_values,
// image_grid_thw,
// pixel_values_video,
// video_grid_thw,
// Some(&cache_position),
// seqlen_offset,
// )?;
// let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
// let next_token = logit_processor.sample(&logits)?;
// let mut decode_ids = Vec::new();
// if !error_tokens.is_empty() {
// decode_ids.extend_from_slice(&error_tokens);
// }
// decode_ids.push(next_token);
// let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?;
// if decoded_token.contains("�") {
// error_tokens.push(next_token);
// if error_tokens.len() > 3 {
// error_tokens.clear();
// }
// seqlen_offset += seq_len;
// seq_len = 1;
// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
// cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
// pixel_values = None;
// pixel_values_video = None;
// continue;
// }
// error_tokens.clear();
// // 处理特殊标记和工具调用
// match decoded_token.as_str() {
// "<tool_call>" => {
// // 开始工具调用
// tool_call_id = Some(uuid::Uuid::new_v4().to_string());
// seqlen_offset += seq_len;
// seq_len = 1;
// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
// cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
// pixel_values = None;
// pixel_values_video = None;
// continue;
// }
// "</tool_call>" => {
// // 结束工具调用
// let chunk = build_completion_chunk_response(
// decoded_token,
// &self.model_name,
// tool_call_id.clone(),
// Some(tool_call_content.clone())
// );
// tool_call_id = None;
// tool_call_content = String::new();
// yield Ok(chunk);
// }
// _ => {
// if tool_call_id.is_some() {
// // 在工具调用过程中，收集工具调用内容
// tool_call_content.push_str(&decoded_token);
// seqlen_offset += seq_len;
// seq_len = 1;
// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
// cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
// pixel_values = None;
// pixel_values_video = None;
// continue;
// } else {
// // 正常文本输出
// let chunk = build_completion_chunk_response(
// decoded_token,
// &self.model_name,
// None,
// None
// );
// yield Ok(chunk);
// }
// }
// }
// if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
// break;
// }
// seqlen_offset += seq_len;
// seq_len = 1;
// input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
// cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
// pixel_values = None;
// pixel_values_video = None;
// }
// self.qwen3_vl.clear_kv_cache();
// };
// Ok(Box::new(Box::pin(stream)))
}
}
src/params/chat.rs
+10
View File
@@ -142,6 +142,16 @@ pub struct ChatCompletionParameters {
pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<usize>,
/// Penalize repetition of token sequences, helping to prevent the model from getting stuck in loops.
/// Values > 1.0 penalize repetitions, values < 1.0 encourage them.
/// For example, a value of 1.2 means that tokens that have already appeared will be less likely to appear again.
#[serde(skip_serializing_if = "Option::is_none")]
pub repeat_penalty: Option<f32>,
/// The number of previous tokens to consider when applying repeat penalty.
/// Only the last N tokens in the conversation history will be used to calculate the penalty.
/// A larger value considers more history for penalizing repetitions, while 0 applies no penalty.
#[serde(skip_serializing_if = "Option::is_none")]
pub repeat_last_n: Option<usize>,
/// A list of tools the model may call. Currently, only functions are supported as a tool.
/// Use this to provide a list of functions the model may generate JSON inputs for.
#[serde(skip_serializing_if = "Option::is_none")]