Skip to content

Commit e89a484

Browse files
authored
[Fix] Fix gemma 3 conv template stop token (#3191)
Currently for all the gemma3 instruction models, we reused the previous gemma conv template. However, there is a difference in the stop token. - In gemma2, the stop tokens are 1, 107: https://huggingface.co/google/gemma-2-2b-it/blob/main/generation_config.json - In gemma3, the stop tokens are 1, 106: https://huggingface.co/google/gemma-3-1b-it/blob/main/generation_config.json This PR add a new `gemma3_instruction` conv template, and will manually update the already-uploaded gemma 3 models on HF. In gemma3, 107 is `\n`. So currently whenever the model tries to change to a new line, it ends generation.
1 parent 2538b87 commit e89a484

File tree

1 file changed

+16
-0
lines changed
  • python/mlc_llm/conversation_template

1 file changed

+16
-0
lines changed

python/mlc_llm/conversation_template/gemma.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,3 +19,19 @@
1919
system_prefix_token_ids=[2],
2020
)
2121
)
22+
23+
# Gemma 3 Instruction. Same as gemma_instruction but with different stop token id
24+
ConvTemplateRegistry.register_conv_template(
25+
Conversation(
26+
name="gemma3_instruction",
27+
system_template=f"{MessagePlaceholders.SYSTEM.value}",
28+
system_message="",
29+
roles={"user": "<start_of_turn>user", "assistant": "<start_of_turn>model"},
30+
seps=["<end_of_turn>\n"],
31+
role_content_sep="\n",
32+
role_empty_sep="\n",
33+
stop_str=["<end_of_turn>"],
34+
stop_token_ids=[1, 106],
35+
system_prefix_token_ids=[2],
36+
)
37+
)

0 commit comments

Comments
 (0)