google-deepmind / google-deepmind/gemma
Bug: `ChatSampler._print_stream` crashes when streaming yields no tokens
- 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`.
Contributor guide
Assessment
This issue has not been assessed yet.