Skip to content

Fix end-of-turn token leaking into streamed chat output - #747

Open
PratikDhanave wants to merge 1 commit into
google-deepmind:mainfrom
PratikDhanaveFork:fix/chat-sampler-stream-end-of-turn-leak
Open

Fix end-of-turn token leaking into streamed chat output#747
PratikDhanave wants to merge 1 commit into
google-deepmind:mainfrom
PratikDhanaveFork:fix/chat-sampler-stream-end-of-turn-leak

Conversation

@PratikDhanave

Copy link
Copy Markdown

What

_print_stream in gemma/gm/text/_chat_sampler.py (used by ChatSampler on the streaming path) is meant to suppress the end-of-turn marker from the streamed/returned text — the comment says "Last token is not printed." But the guard is placed after the work it should gate:

for state in out:
    print_(stream, state.text)          # printed before the check
    text_tokens.append(state.text)      # appended before the check
    if state.text == '<end_of_turn>' or state.text == '<turn|>':
        continue                        # no-op: last statement in the loop body
out = dataclasses.replace(state, text=''.join(text_tokens))

Because print_ and append run before the if, and continue is the final statement in the loop body (so it does nothing), the end-of-turn marker (<end_of_turn> / <turn|>) is:

  1. streamed to the user, and
  2. concatenated into the returned SamplerOutput.text.

Any streamed chat turn ending normally therefore leaks the literal end-of-turn token into the user-visible output.

Change

Move the guard ahead of the print/append so the marker is skipped, matching the documented intent. No other behavior change.

for state in out:
    if state.text == '<end_of_turn>' or state.text == '<turn|>':
        # Last token is not printed.
        continue
    print_(stream, state.text)
    text_tokens.append(state.text)

In `_print_stream`, the guard meant to suppress the end-of-turn marker
(`<end_of_turn>` / `<turn|>`) 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant