From 62c7e8f932df2829581ab987ff6d0322945b5f50 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 5 Nov 2026 19:53:10 +0900 Subject: [PATCH] fix(generate): start a chat reply on its own line after the prompt echo `mlxcel generate` echoes the user prompt and then streams the reply with no separator, so a chat-templated reply runs on from the prompt ("...on a laptop.Cores ablaze bright"). With a chat template the reply is a new assistant turn, so the echo now ends with a newline. Without one (`--no-chat-template`, or a checkpoint that has no template) the reply continues the prompt text and the echo is unchanged. This matches llama.cpp: `llama-cli` echoes the prompt as `> prompt` on its own line before the reply, `llama-completion` in conversation mode echoes the rendered template so the reply follows the assistant header, and only raw `-no-cnv` completion prints prompt and continuation as one passage. The text, pipeline, and XLA paths decide from whether the rendered prompt differs from the user prompt; the diffusion, LLaDA-2, and Florence-2 paths from `--no-chat-template`; Nemotron Parse always starts its output on a new line. Validated with two unit tests and real runs on an M1 Ultra: Meta-Llama-3.1-8B-Instruct-4bit and gemma-4-e4b-it-4bit put the reply on the line after the prompt, and `--no-chat-template -p "The capital of France is"` still prints "The capital of France is Paris.". --- src/commands/generate.rs | 28 ++++++++++++++++++++----- src/commands/generate_diffusion.rs | 2 +- src/commands/generate_florence2.rs | 2 +- src/commands/generate_llada2.rs | 2 +- src/commands/generate_nemotron_parse.rs | 2 +- src/commands/generate_tests.rs | 24 +++++++++++++++++++++ 6 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 5f1433f72..aa43524eb 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1270,13 +1270,31 @@ fn build_cli_sampling_config_with_flags( build_sampling_config(resolved_cli_sampling_params(args, stop_token_ids, flags)) } -pub(super) fn print_generation_preamble(user_prompt: &str) -> Result<()> { +/// Echo the prompt before the reply streams in. +/// +/// With a chat template the reply is a new assistant turn, so the echoed prompt +/// ends with a newline and the reply starts on its own line. Without one (raw +/// completion, `--no-chat-template`, or a checkpoint without a template) the +/// reply continues the prompt text, so no newline is added and the prompt and +/// completion read as one passage. +pub(super) fn print_generation_preamble(user_prompt: &str, reply_is_new_turn: bool) -> Result<()> { println!("Generating..."); - print!("{}", user_prompt); + print!( + "{}", + generation_preamble_echo(user_prompt, reply_is_new_turn) + ); io::stdout().flush()?; Ok(()) } +fn generation_preamble_echo(user_prompt: &str, reply_is_new_turn: bool) -> String { + if reply_is_new_turn && !user_prompt.ends_with('\n') { + format!("{user_prompt}\n") + } else { + user_prompt.to_string() + } +} + fn generated_suffix<'a>(full_text: &'a str, prompt_text: &str) -> &'a str { full_text.strip_prefix(prompt_text).unwrap_or(full_text) } @@ -3207,7 +3225,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { if select_backend().name() == "xla" { validate_xla_output_audio(args.generation.output_audio.as_deref())?; let num_layers = xla_num_layers(&args.model.model); - print_generation_preamble(&user_prompt)?; + print_generation_preamble(&user_prompt, prompt != user_prompt)?; let (generated_tokens, stats) = generate_xla( &args.model.model, num_layers, @@ -3254,7 +3272,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { &[], ); let num_layers = resolve_cli_pipeline_num_layers(&args.model.model)?; - print_generation_preamble(&user_prompt)?; + print_generation_preamble(&user_prompt, prompt != user_prompt)?; generate_pipeline_text( &args.model.model, num_layers, @@ -3363,7 +3381,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { no_chat_template: args.generation.no_chat_template, }, )?; - print_generation_preamble(&user_prompt)?; + print_generation_preamble(&user_prompt, prompt != user_prompt)?; let generation = run_generation_mode( &model, &args, diff --git a/src/commands/generate_diffusion.rs b/src/commands/generate_diffusion.rs index f82971720..115f7812f 100644 --- a/src/commands/generate_diffusion.rs +++ b/src/commands/generate_diffusion.rs @@ -176,7 +176,7 @@ pub(crate) fn run_diffusion_generation( None => (prompt_tokens, None), }; - print_generation_preamble(user_prompt)?; + print_generation_preamble(user_prompt, !args.generation.no_chat_template)?; // Stream through the shared incremental detokenizer (byte-fallback // safe); raw ids are collected in parallel so any held-back tail bytes diff --git a/src/commands/generate_florence2.rs b/src/commands/generate_florence2.rs index 9a7bb394d..decbd88a9 100644 --- a/src/commands/generate_florence2.rs +++ b/src/commands/generate_florence2.rs @@ -79,7 +79,7 @@ pub(crate) fn run_florence2_generation( .pop() .ok_or_else(|| anyhow!("Image decoding returned no image for {image_path:?}"))?; - print_generation_preamble(user_prompt)?; + print_generation_preamble(user_prompt, !args.generation.no_chat_template)?; println!(); let started = Instant::now(); diff --git a/src/commands/generate_llada2.rs b/src/commands/generate_llada2.rs index e4c9aea6e..76a2f7b69 100644 --- a/src/commands/generate_llada2.rs +++ b/src/commands/generate_llada2.rs @@ -100,7 +100,7 @@ pub(crate) fn run_llada2_generation( mlxcel_core::random_seed(seed); } - print_generation_preamble(user_prompt)?; + print_generation_preamble(user_prompt, !args.generation.no_chat_template)?; // Stream through the shared incremental detokenizer; raw ids are collected // in parallel so any held-back tail bytes can be flushed afterward. diff --git a/src/commands/generate_nemotron_parse.rs b/src/commands/generate_nemotron_parse.rs index b81d4fde1..a53eff838 100644 --- a/src/commands/generate_nemotron_parse.rs +++ b/src/commands/generate_nemotron_parse.rs @@ -85,7 +85,7 @@ pub(crate) fn run_nemotron_parse_generation( .pop() .ok_or_else(|| anyhow!("Image decoding returned no image for {image_path:?}"))?; - print_generation_preamble(prompt)?; + print_generation_preamble(prompt, true)?; println!(); let started = Instant::now(); diff --git a/src/commands/generate_tests.rs b/src/commands/generate_tests.rs index 9762511f2..9b392bfaa 100644 --- a/src/commands/generate_tests.rs +++ b/src/commands/generate_tests.rs @@ -14,6 +14,7 @@ #[cfg(feature = "xla-backend")] use super::decode_xla_cli_images; +use super::generation_preamble_echo; use super::{ CliPromptMedia, CliSamplingFlagState, CliVideoFrameGroup, CliVideoFrames, OfflineDraftRoute, apply_user_chat_template, apply_vlm_chat_template, build_cli_sampling_config_with_flags, @@ -1707,3 +1708,26 @@ fn the_repl_uses_a_drafter_only_with_an_explicit_server_kind() { args.speculative.draft_kind = Some("mtp".to_string()); assert_eq!(repl_drafter(&args), None, "a kind without a drafter"); } + +#[test] +fn preamble_echo_ends_the_prompt_line_before_a_chat_reply() { + // A chat-templated reply is a new assistant turn: it must not run on from + // the echoed prompt ("...on a laptop.Cores ablaze"). + assert_eq!( + generation_preamble_echo("Write a haiku.", true), + "Write a haiku.\n" + ); + assert_eq!( + generation_preamble_echo("Already ends\n", true), + "Already ends\n" + ); +} + +#[test] +fn preamble_echo_lets_a_raw_completion_continue_the_prompt() { + // Without a chat template the completion continues the prompt text. + assert_eq!( + generation_preamble_echo("The capital of France is", false), + "The capital of France is" + ); +}