Fix ChatSampler.conversation crash for non-Gemma4 tokenizers (breaks ToolSampler multi-turn)

[sohumt123]

Problem

ChatSampler.conversation raises ValueError for all non-Gemma4 models after any successful chat() turn, and this crashes ToolSampler multi-turn and tool-call flows for those model families.

ChatSampler stores each turn's text pre-formatted with self.tokenizer.FORMAT — which is dialog.Format.GEMMA3 (<start_of_turn>...<end_of_turn> tags) for the Gemma 2/3/3n tokenizers. But the conversation property passes the joined text directly to dialog.Conversation(str), which only parses the canonical Gemma4 <|turn>...<turn|> tags:

ValueError: Conversation root must only contain `turn` tags. Got
'<start_of_turn>user\nHello world<end_of_turn>\n<start_of_turn>model\nhi there<end_of_turn>'

The impact goes beyond the public conversation property: ToolSampler.chat evaluates self.conversation on every call (if not self.conversation or not multi_turn, _tool_sampler.py:107), including the internal recursive call made after a tool call executes. So for Gemma 2/3/3n models, turn 1 succeeds but every second turn raises, fully breaking tool use and multi-turn tool chat.

The mismatch was introduced when the turns storage switched to holding text pre-formatted with the tokenizer's legacy format while conversation migrated to the dialog library's parser, with no format conversion between the two.

Fix

In the conversation property, convert the stored text to the canonical format before parsing:

text = ''.join(t.text for t in self.turns)
if self.tokenizer.FORMAT is not dialog.Format.GEMMA4:
  text = self.tokenizer.FORMAT.to_gemma4(text)
return dialog.Conversation(text)

This is a no-op for Gemma4 tokenizers (FORMAT is GEMMA4) and covers all current dialog.Format members.

Testing

Added two regression tests in gemma/gm/text/_sampler_test.py, following the existing mock-based ChatSampler test pattern (gm.testing.DummyGemma + DummyTokenizer, CPU-only, no downloads):

• test_chat_sampler_conversation_legacy_format — after a chat() turn with a GEMMA3-format tokenizer, conversation returns a dialog.Conversation with the expected User/Model turns. Fails with the ValueError above without the fix.
• test_tool_sampler_multi_turn_legacy_format — ToolSampler.chat completes a second turn with a legacy-format tokenizer. Fails on turn 2 without the fix.

$ pytest gemma/gm/text/ -q
15 passed in 10.99s

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.