perf: stop generation on default turn boundaries

This commit is contained in:
Evandro Fonseca Junior committed 2026-03-26 17:52:59 -03:00
1 parent 1bd3365a31
commit 653f5b5bc1
2 files changed
+81 -12

No files matched your search

+27 -3
View File
@@ -514,7 +514,17 @@ pub fn build_sampler(req: &ChatCompletionRequest) -> inference_core::SamplerConf
sampler.temperature = req.temperature.clamp(0.0, 2.0);
sampler.top_p = req.top_p.clamp(0.0, 1.0);
sampler.max_tokens = req.max_tokens.max(1);
sampler.stop_sequences = req.stop.clone();
sampler.stop_sequences = if req.stop.is_empty() {
vec![
"\nUser:".to_owned(),
"\nSystem:".to_owned(),
"\nDeveloper:".to_owned(),
"\nAssistant:".to_owned(),
"\nTool:".to_owned(),
]
} else {
req.stop.clone()
};
sampler.repeat_penalty = if req.tools.is_empty() { 1.05 } else { 1.0 };
sampler.seed = req.seed.unwrap_or(0);
sampler
@@ -565,8 +575,8 @@ pub fn ensure_supported_mode(req: &ChatCompletionRequest) -> Result<()> {
#[cfg(test)]
mod tests {
use super::{
authorize_request, finalize_assistant_message, prepare_prompt, validate_request,
FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode,
authorize_request, build_sampler, finalize_assistant_message, prepare_prompt,
validate_request, FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode,
};
use api_contracts::chat::ChatCompletionRequest;
use axum::http::{HeaderMap, HeaderValue, StatusCode};
@@ -718,4 +728,18 @@ mod tests {
Some("Use references and lifetimes.")
);
}
#[test]
fn adds_default_turn_stops_when_request_has_none() {
let request = parse_request(
r#"{
"model": "alias-model",
"messages": [{"role":"user","content":"hi"}]
}"#,
);
let sampler = build_sampler(&request);
assert!(sampler.stop_sequences.contains(&"\nUser:".to_owned()));
assert!(sampler.stop_sequences.contains(&"\nAssistant:".to_owned()));
}
}