From b71a564184d36791d0efe28c84e4715785c701ab Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Sat, 25 Jul 2026 10:26:48 +0530 Subject: [PATCH] Fix end-of-turn token leaking into streamed chat output In `_print_stream`, the guard meant to suppress the end-of-turn marker (`` / ``) was placed *after* the print and append, with `continue` as the loop body's final statement -- making it a no-op. As a result the end-of-turn marker was both streamed to the user and concatenated into the returned `SamplerOutput.text`, contradicting the "Last token is not printed" comment. Move the guard ahead of the print/append so the marker is skipped as intended. --- gemma/gm/text/_chat_sampler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index 887ee2ad..df730f0b 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/gemma/gm/text/_chat_sampler.py @@ -468,13 +468,13 @@ def _print_stream( text_tokens = [] for state in out: - print_(stream, state.text) # pyrefly: ignore[bad-argument-type] - - text_tokens.append(state.text) if ( state.text == '' or state.text == '' ): # Last token is not printed. continue + print_(stream, state.text) # pyrefly: ignore[bad-argument-type] + + text_tokens.append(state.text) out = dataclasses.replace(state, text=''.join(text_tokens)) # pylint: disable=undefined-variable,undefined-loop-variable # pyrefly: ignore[bad-assignment] return out # pyrefly: ignore[bad-return]