google-deepmind / google-deepmind/gemma

Bug: `ChatSampler._print_stream` crashes when streaming yields no tokens

Open
#551 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
5.7k
Forks
1k
Avg merge
10h 33m
Merged PRs (30d)
2

Description

### Summary

When using `gm.text.ChatSampler` with `print_stream=True`, the internal helper function `_print_stream` assumes that the streaming iterator always yields at least one `SamplerOutput`. If the iterator is empty (for example, due to `max_new_tokens=0` or an early stop condition before the first token is emitted), `_print_stream` attempts to access the loop variable `state` even though the loop body was never executed. This results in a runtime error instead of a graceful empty response.

### Affected code

- File: `gemma/gm/text/_chat_sampler.py`
- Function: `_print_stream`

Relevant snippet:

```python
def _print_stream(
out: Iterator[_sampler.SamplerOutput],
) -> _sampler.SamplerOutput:
"""Prints the streaming output."""
text_tokens = []
for state in out:
text_tokens.append(state.text)
if state.text == '': # Last token is not printed.
continue
print(state.text, end='', flush=True)
out = dataclasses.replace(state, text=''.join(text_tokens)) # pylint: disable=undefined-variable,undefined-loop-variable
return out
```

### Steps to reproduce

1. Ensure `gemma` is installed and importable.
2. Create a minimal script that calls `_print_stream` with an empty iterator. For example:

```python
from gemma.gm.text import _chat_sampler

def main() -> None:
empty_iterator = iter(())
_chat_sampler._print_stream(empty_iterator)

if __name__ == "__main__":
main()
```

3. Run the script:

```bash
python examples/demo_chat_sampler_print_stream_nameerror.py
```

### Actual behavior

The script crashes with an internal error similar to:

```text
File "gemma/gm/text/_chat_sampler.py", line 233, in _print_stream
out = dataclasses.replace(state, text=''.join(text_tokens))
UnboundLocalError: cannot access local variable 'state' where it is not associated with a value
```

This happens because the `for state in out:` loop never executes for an empty iterator, so `state` is never bound before it is used in `dataclasses.replace`.

### Expected behavior

`ChatSampler.chat(..., print_stream=True)` should handle the case where the streaming sampler yields no tokens **without** crashing. In such cases, it should either:

- Return a valid `SamplerOutput` with an empty `text` field, or
- Raise a clear, high-level error message describing that no tokens were generated,

but it should not fail with an internal `UnboundLocalError` inside `_print_stream`.

Image

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.