mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-19 07:48:29 +08:00
Make Gemma 4 channel close markers channel aware
decode() rewrote every <channel|> into </think>, including the close of a channel that is not the thought channel. Only close a thought channel that was actually opened in the decoded text; other channel markers are left alone for the caller to strip.
This commit is contained in:
@@ -6,6 +6,7 @@ import numpy as np
|
||||
from tokenizers import Tokenizer
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
import re
|
||||
|
||||
from comfy import sd1_clip
|
||||
import comfy.model_management
|
||||
@@ -1401,9 +1402,11 @@ class Gemma4SDTokenizer(Gemma4_Tokenizer, sd1_clip.SDTokenizer):
|
||||
|
||||
def decode(self, token_ids, **kwargs):
|
||||
text = super().decode(token_ids, skip_special_tokens=False)
|
||||
# Translate thinking channel markers to standard <think>/</think> tags
|
||||
# Translate thinking channel markers to standard <think>/</think> tags. Only a close
|
||||
# that ends a thought channel becomes </think>; generation primed with another channel
|
||||
# leaves its opening marker in the prompt, so its close is not reasoning.
|
||||
text = re.sub(r"<\|channel>thought\n(.*?)<channel\|>", r"<think>\n\1</think>", text, flags=re.DOTALL)
|
||||
text = text.replace("<|channel>thought\n", "<think>\n")
|
||||
text = text.replace("<channel|>", "</think>")
|
||||
# Strip remaining special tokens
|
||||
text = text.replace("<turn|>", "").replace("<eos>", "").strip()
|
||||
return text
|
||||
|
||||
45
tests-unit/comfy_test/test_gemma4_tokenizer.py
Normal file
45
tests-unit/comfy_test/test_gemma4_tokenizer.py
Normal file
@@ -0,0 +1,45 @@
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
from comfy.text_encoders.gemma4 import Gemma4SDTokenizer # noqa: E402
|
||||
|
||||
|
||||
class _StubTokenizer:
|
||||
"""Returns a canned decode so the marker translation can be tested without model files."""
|
||||
def __init__(self, text):
|
||||
self.text = text
|
||||
|
||||
def decode(self, token_ids, skip_special_tokens=False):
|
||||
return self.text
|
||||
|
||||
|
||||
def decode(text):
|
||||
tokenizer = Gemma4SDTokenizer.__new__(Gemma4SDTokenizer)
|
||||
tokenizer.tokenizer = _StubTokenizer(text)
|
||||
return tokenizer.decode([])
|
||||
|
||||
|
||||
class TestGemma4Decode:
|
||||
def test_thought_channel_becomes_think_tags(self):
|
||||
assert decode("<|channel>thought\nreasoning<channel|>the answer") == "<think>\nreasoning</think>the answer"
|
||||
|
||||
def test_primed_empty_thought_channel(self):
|
||||
assert decode("<|channel>thought\n<channel|>the answer") == "<think>\n</think>the answer"
|
||||
|
||||
def test_unclosed_thought_channel(self):
|
||||
assert decode("<|channel>thought\nreasoning") == "<think>\nreasoning"
|
||||
|
||||
def test_other_channel_close_is_not_reasoning(self):
|
||||
# Non-thinking LTX2 prompt enhancement primes a "final" channel, so only its close is
|
||||
# generated. Turning that into </think> made the whole answer look like reasoning.
|
||||
assert decode("the answer<channel|>") == "the answer<channel|>"
|
||||
|
||||
def test_other_channel_kept_after_a_thought_channel(self):
|
||||
assert decode("<|channel>thought\nreasoning<channel|><|channel>final\nthe answer<channel|>") == "<think>\nreasoning</think><|channel>final\nthe answer<channel|>"
|
||||
|
||||
def test_turn_and_eos_are_stripped(self):
|
||||
assert decode("the answer<turn|><eos>") == "the answer"
|
||||
Reference in New Issue
Block a user