diff --git a/crates/inference-core/src/session.rs b/crates/inference-core/src/session.rs index a0defd6..02dd0af 100644 --- a/crates/inference-core/src/session.rs +++ b/crates/inference-core/src/session.rs @@ -27,6 +27,8 @@ pub struct GenerationSession { vocab: *const llama_vocab, sampler: *mut llama_sys::llama_sampler, stop_sequences: Vec, + stop_suffix_len: usize, + generated_tail: String, // generation state n_decoded: i32, // number of tokens decoded so far (prefill + gen) @@ -163,12 +165,21 @@ impl GenerationSession { ); // create session state + let stop_suffix_len = sampler_cfg + .stop_sequences + .iter() + .map(|stop| stop.len()) + .max() + .unwrap_or(0); + Ok(Self { lib, ctx, vocab, sampler: sampler_chain, // use sampler_chain here stop_sequences: sampler_cfg.stop_sequences.clone(), + stop_suffix_len, + generated_tail: String::new(), n_decoded: n_prompt_tokens, prompt_tokens: n_prompt_tokens as u32, max_tokens: sampler_cfg.max_tokens, @@ -227,14 +238,10 @@ impl GenerationSession { let token_text = self.detokenize_token(new_token); // Check stop sequences - // (simple check: if generated text so far ends with a stop seq) - // For streaming we check per-token — good enough for Phase 2. - for stop in &self.stop_sequences { - if token_text.trim_end() == stop.as_str() { - self.finished = true; - self.finalize_timing(); - return None; - } + if self.matches_stop_sequence(&token_text) { + self.finished = true; + self.finalize_timing(); + return None; } Some(token_text) @@ -291,6 +298,28 @@ impl GenerationSession { .into_owned() } } + + fn matches_stop_sequence(&mut self, token_text: &str) -> bool { + if self.stop_sequences.is_empty() { + return false; + } + + self.generated_tail.push_str(token_text); + tail_ends_with_stop( + &mut self.generated_tail, + self.stop_suffix_len, + &self.stop_sequences, + ) + } +} + +fn tail_ends_with_stop(tail: &mut String, max_len: usize, stop_sequences: &[String]) -> bool { + if tail.len() > max_len { + let trim_from = tail.len() - max_len; + tail.drain(..trim_from); + } + + stop_sequences.iter().any(|stop| tail.ends_with(stop)) } fn decode_token_piece_with(mut call: F) -> String @@ -318,7 +347,7 @@ where #[cfg(test)] mod tests { - use super::decode_token_piece_with; + use super::{decode_token_piece_with, tail_ends_with_stop}; #[test] fn retries_when_piece_buffer_is_too_small() { @@ -364,6 +393,22 @@ mod tests { assert_eq!(decoded, " Quantum computer"); assert_eq!(calls, 2); } + + #[test] + fn matches_stop_sequence_across_multiple_tokens() { + let stops = vec!["\nAssistant:".to_owned()]; + let mut tail = String::new(); + + tail.push_str("\nAssis"); + assert!(!tail_ends_with_stop( + &mut tail, + "\nAssistant:".len(), + &stops + )); + + tail.push_str("tant:"); + assert!(tail_ends_with_stop(&mut tail, "\nAssistant:".len(), &stops)); + } } impl Drop for GenerationSession { diff --git a/crates/inference-server/src/compat.rs b/crates/inference-server/src/compat.rs index 0adbc6d..5e1eff9 100644 --- a/crates/inference-server/src/compat.rs +++ b/crates/inference-server/src/compat.rs @@ -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())); + } }