wan/xgrammar-gemma4-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
wanhae.lee
2026-07-05 00:33:50 +09:00
commit 9484167719
7 changed files with 1125 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+14
View File
@@ -0,0 +1,14 @@
# Python
__pycache__/
*.py[cod]
*.egg-info/
# Virtual environments
.venv/
venv/
# uv
uv.lock
# OS
.DS_Store
Makefile
+37
View File
@@ -0,0 +1,37 @@
XGRAMMAR_DIR ?= $(HOME)/git/xgrammar
VENV := .venv
PY := $(VENV)/bin/python
SAMPLES ?= 4
TEMPERATURE ?= 0.9
.PHONY: run smoke setup xgrammar clean help
help:
@echo "make setup - create uv venv and install deps + local xgrammar (editable)"
@echo "make run - run all scenarios and write report.md (constrained vs unconstrained)"
@echo "make smoke - quick single-tool smoke test"
@echo "make xgrammar - reinstall local xgrammar after C++/python changes"
@echo "make clean - remove venv and report"
$(VENV)/.stamp:
uv venv --python 3.13 $(VENV)
uv pip install --python $(PY) torch transformers accelerate jsonschema
uv pip install --python $(PY) -e $(XGRAMMAR_DIR)
touch $@
setup: $(VENV)/.stamp
xgrammar: $(VENV)/.stamp
uv pip install --python $(PY) -e $(XGRAMMAR_DIR)
run: $(VENV)/.stamp
$(PY) test_gemma4_scenarios.py --samples $(SAMPLES) --temperature $(TEMPERATURE) --report report.md
@echo ""
@echo "==================== report.md ===================="
@cat report.md
smoke: $(VENV)/.stamp
$(PY) test_gemma4_xgrammar.py
clean:
rm -rf $(VENV) report.md
README.md
+65
View File
@@ -0,0 +1,65 @@
# gemma-4 × xgrammar constrained decoding test
A collection of scripts that verify the modified xgrammar (`~/git/xgrammar`,
`style="gemma"` / `gemma_4` builtin structural tag) works correctly for tool
calling with `google/gemma-4-E2B-it`, and generate a report comparing
constrained vs unconstrained behavior.
## Usage (Makefile)
```bash
make setup # create uv venv + install deps + local xgrammar editable install (once)
make run # run all scenarios → generate report.md + print output ← for PR reports
make smoke # single-tool smoke test (greedy, fast)
make xgrammar # reinstall after modifying ~/git/xgrammar
make clean # remove venv/report
```
Options: `make run SAMPLES=8`, `XGRAMMAR_DIR=/path/to/xgrammar make setup`
The output of `make run`, **`report.md`**, is the artifact to show reviewers —
it contains the model/sampling configuration, a per-scenario table of
constrained vs unconstrained metrics, and the raw text of failure cases.
## Scenarios (`test_gemma4_scenarios.py`)
| | Setup | What it targets |
|---|---|---|
| **A. complex-schema** | Nested objects + array of objects + enum + boolean schema, titles containing quotes | Argument serialization mistakes |
| **B. multiturn-thinking** | Two prior rounds of tool call/response already in context and thinking is required, the next call takes a number-typed argument | Skipped/unterminated thinking, type errors |
| **C. adversarial-payload** | A string argument whose body must contain JSON/braces/an error log verbatim, an out-of-enum word ("critical") lure, two similarly named distractor tools, temp 1.5 + top_p 1.0 | String quoting collapse, enum violations |
Metrics (pass/fail per sample):
| metric | meaning |
|---|---|
| `well_formed` | Every `<\|tool_call>call:name{...}<tool_call\|>` block is complete and parseable |
| `valid_name` | Only calls tools that actually exist |
| `schema_valid` | Parsed arguments pass the JSON schema (parsed via a port of the sglang parser) |
| `stops_at_boundary` | Generation doesn't run past `<tool_call\|>` into `<\|tool_response>` |
| `thought_ok` | (thinking scenario) Opens the thought channel and closes it before the tool call |
## Key findings
- gemma-4-E2B-it's format training is very robust, so the tool call argument
format itself rarely breaks even at temp 1.5 + top_p 1.0.
- Per the [Gemma 4 prompt-formatting spec](https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4),
the model is only responsible for generating up to `<tool_call\|>`;
`<\|tool_response>` is appended by the application with the real tool
result, and is registered as an **additional stop sequence** purely as a
backstop. This harness has no such stop configured, so unconstrained runs
keep generating past that boundary into engine-owned territory
(`stops_at_boundary` failures); required-mode constrained decoding ends the
call cleanly at an accept state instead.
- Unconstrained decoding also **omits the thought channel entirely** in
situations where thinking should be enabled; constrained decoding with
`reasoning=True` enforces it (an empty thought `<\|channel>thought\n<channel\|>`
remains legal, matching the spec's no-thinking form).
## Files
- `test_gemma4_xgrammar.py` — basic smoke test (greedy, single tool, alignment check)
- `test_gemma4_scenarios.py` — scenario runner + report.md generator.
`--scenario a|b|c|all --samples N --temperature T --model ID --report PATH`
- `gemma_parser.py` — a pure port of sglang's `Gemma4Detector` parsing logic (for output verification)
- `Makefile` / `report.md`
gemma_parser.py
+215
View File
@@ -0,0 +1,215 @@
"""Minimal Gemma4 tool-call parser.
Pure-function port of sglang's Gemma4Detector parsing helpers
(sglang/python/sglang/srt/function_call/gemma4_detector.py), so generated
outputs can be validated exactly the way sglang would parse them, without
importing sglang.
"""
TOOL_CALL_START = "<|tool_call>"
TOOL_CALL_END = "<tool_call|>"
STRING_DELIM = '<|"|>'
def _parse_value(value_str: str) -> object:
value_str = value_str.strip()
if not value_str:
return value_str
if value_str == "true":
return True
if value_str == "false":
return False
try:
if "." in value_str:
return float(value_str)
return int(value_str)
except ValueError:
pass
return value_str
def _parse_array(arr_str: str) -> list:
items: list = []
i = 0
n = len(arr_str)
while i < n:
while i < n and arr_str[i] in (" ", ",", "\n", "\t"):
i += 1
if i >= n:
break
if arr_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
i += len(STRING_DELIM)
end_pos = arr_str.find(STRING_DELIM, i)
if end_pos == -1:
items.append(arr_str[i:])
break
items.append(arr_str[i:end_pos])
i = end_pos + len(STRING_DELIM)
elif arr_str[i] == "{":
depth = 1
obj_start = i + 1
i += 1
while i < n and depth > 0:
if arr_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
i += len(STRING_DELIM)
next_delim = arr_str.find(STRING_DELIM, i)
i = next_delim + len(STRING_DELIM) if next_delim != -1 else n
continue
if arr_str[i] == "{":
depth += 1
elif arr_str[i] == "}":
depth -= 1
i += 1
items.append(parse_args(arr_str[obj_start : i - 1]))
elif arr_str[i] == "[":
depth = 1
sub_start = i + 1
i += 1
while i < n and depth > 0:
if arr_str[i] == "[":
depth += 1
elif arr_str[i] == "]":
depth -= 1
i += 1
items.append(_parse_array(arr_str[sub_start : i - 1]))
else:
val_start = i
while i < n and arr_str[i] not in (",", "]"):
i += 1
items.append(_parse_value(arr_str[val_start:i]))
return items
def parse_args(args_str: str) -> dict:
"""Parse Gemma4's key:value argument format into a Python dict."""
if not args_str or not args_str.strip():
return {}
result: dict = {}
i = 0
n = len(args_str)
while i < n:
while i < n and args_str[i] in (" ", ",", "\n", "\t"):
i += 1
if i >= n:
break
key_start = i
while i < n and args_str[i] != ":":
i += 1
if i >= n:
break
key = args_str[key_start:i].strip()
i += 1
if i >= n:
result[key] = ""
break
while i < n and args_str[i] in (" ", "\n", "\t"):
i += 1
if i >= n:
result[key] = ""
break
if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
i += len(STRING_DELIM)
val_start = i
end_pos = args_str.find(STRING_DELIM, i)
if end_pos == -1:
result[key] = args_str[val_start:]
break
result[key] = args_str[val_start:end_pos]
i = end_pos + len(STRING_DELIM)
elif args_str[i] == "{":
depth = 1
obj_start = i + 1
i += 1
while i < n and depth > 0:
if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
i += len(STRING_DELIM)
next_delim = args_str.find(STRING_DELIM, i)
i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
continue
if args_str[i] == "{":
depth += 1
elif args_str[i] == "}":
depth -= 1
i += 1
result[key] = parse_args(args_str[obj_start : i - 1])
elif args_str[i] == "[":
depth = 1
arr_start = i + 1
i += 1
while i < n and depth > 0:
if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
i += len(STRING_DELIM)
next_delim = args_str.find(STRING_DELIM, i)
i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
continue
if args_str[i] == "[":
depth += 1
elif args_str[i] == "]":
depth -= 1
i += 1
result[key] = _parse_array(args_str[arr_start : i - 1])
else:
val_start = i
while i < n and args_str[i] not in (",", "}", "]"):
i += 1
result[key] = _parse_value(args_str[val_start:i])
return result
def _find_matching_brace(text: str) -> int:
"""Index of the matching '}' in text (which starts just after '{'), or -1."""
depth = 1
i = 0
n = len(text)
dlen = len(STRING_DELIM)
while i < n and depth > 0:
if text[i : i + dlen] == STRING_DELIM:
i += dlen
next_delim = text.find(STRING_DELIM, i)
if next_delim == -1:
return -1
i = next_delim + dlen
continue
if text[i] == "{":
depth += 1
elif text[i] == "}":
depth -= 1
i += 1
return (i - 1) if depth == 0 else -1
def extract_tool_calls(text: str) -> list:
"""Extract [(func_name, args_dict_or_None)] from text.
args is None when the call block is malformed (unparseable), mirroring
where sglang's detector would fail.
"""
results = []
search_from = 0
while True:
start = text.find(TOOL_CALL_START, search_from)
if start == -1:
break
end = text.find(TOOL_CALL_END, start)
if end == -1:
# unterminated call block
results.append((None, None))
break
inner = text[start + len(TOOL_CALL_START) : end]
if inner.startswith("call:"):
brace = inner.find("{")
if brace != -1:
func_name = inner[5:brace]
args_content = inner[brace + 1 :]
match_idx = _find_matching_brace(args_content)
args_str = args_content[:match_idx] if match_idx != -1 else args_content
try:
results.append((func_name, parse_args(args_str)))
except Exception:
results.append((func_name, None))
else:
results.append((None, None))
else:
results.append((None, None))
search_from = end + len(TOOL_CALL_END)
return results
report.md
+128
View File
@@ -0,0 +1,128 @@
# Gemma-4 tool calling: constrained (xgrammar) vs unconstrained
- model: `google/gemma-4-E2B-it` via HF transformers
- constraint: `xgr.get_model_structural_tag("gemma_4", tools=..., tool_choice="required")` compiled and applied with `xgr.contrib.hf.LogitsProcessor`
- 4 samples per scenario per mode; identical sampling settings for both modes (per-scenario values below)
- outputs validated with a pure-Python port of sglang's `Gemma4Detector` parser plus `jsonschema`
## Results (passing samples / total)
### A. complex-schema
temperature=0.9, top_p=0.95, max_new_tokens=384, reasoning=False
| metric | constrained | unconstrained |
|---|---|---|
| well_formed | 4/4 ✅ | 4/4 ✅ |
| valid_name | 4/4 ✅ | 4/4 ✅ |
| schema_valid | 4/4 ✅ | 4/4 ✅ |
| stops_at_boundary | 4/4 ✅ | 0/4 ❌ |
### B. multiturn-thinking
temperature=0.9, top_p=0.95, max_new_tokens=384, reasoning=True
| metric | constrained | unconstrained |
|---|---|---|
| well_formed | 4/4 ✅ | 4/4 ✅ |
| valid_name | 4/4 ✅ | 4/4 ✅ |
| schema_valid | 4/4 ✅ | 4/4 ✅ |
| stops_at_boundary | 4/4 ✅ | 0/4 ❌ |
| thought_ok | 4/4 ✅ | 0/4 ❌ |
### C. adversarial-payload
temperature=1.5, top_p=1.0, max_new_tokens=640, reasoning=False
| metric | constrained | unconstrained |
|---|---|---|
| well_formed | 4/4 ✅ | 4/4 ✅ |
| valid_name | 4/4 ✅ | 4/4 ✅ |
| schema_valid | 4/4 ✅ | 4/4 ✅ |
| stops_at_boundary | 4/4 ✅ | 0/4 ❌ |
## Metric definitions
- `well_formed` — every tool call block is complete and parseable (sglang parser port)
- `valid_name` — every called tool exists
- `schema_valid` — parsed arguments validate against the tool's JSON schema
- `stops_at_boundary` — generation did not run past <tool_call|> into <|tool_response>
- `thought_ok` — <|channel>thought opened and closed before the first tool call
## Failure examples
**A. complex-schema / unconstrained #0** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
```
**A. complex-schema / unconstrained #1** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
```
**A. complex-schema / unconstrained #2** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
```
**A. complex-schema / unconstrained #3** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
```
**B. multiturn-thinking / unconstrained #0** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
```
<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
```
**B. multiturn-thinking / unconstrained #1** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
```
<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
```
**B. multiturn-thinking / unconstrained #2** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
```
<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
```
**B. multiturn-thinking / unconstrained #3** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
```
<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
```
**C. adversarial-payload / unconstrained #0** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42).<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
```
**C. adversarial-payload / unconstrained #1** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42).<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>bug<|"|>,<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
```
**C. adversarial-payload / unconstrained #2** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42)<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>bug<|"|>,<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
```
**C. adversarial-payload / unconstrained #3** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
```
<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42).<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>bug<|"|>,<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
```
test_gemma4_scenarios.py
+546
View File
@@ -0,0 +1,546 @@
"""Constrained-vs-unconstrained tool calling scenarios for Gemma-4 + xgrammar.
Compares grammar-constrained generation (xgrammar gemma_4 builtin structural
tag, JSONSchemaFormat style="gemma") against unconstrained generation on
google/gemma-4-E2B-it, across scenarios designed to break tool call syntax:
A. complex-schema — nested objects, array of objects, enum, boolean,
integer; prompt values bait JSON-style quoting.
B. multiturn-thinking — two prior tool call/response rounds in context,
thinking enabled; the model must close a
<|channel>thought section, then call a *different*
tool whose `value` argument must be a number.
C. adversarial-payload — a string argument that must contain a JSON snippet
full of braces/quotes (breaks brace-matching if the
model quotes it wrong), an enum baited with a word
that is NOT in the enum ("critical"), and two
similarly-named distractor tools. Runs at elevated
temperature (1.3) to expose format instability.
Metrics per sample:
well_formed — every <|tool_call> block is complete and parseable by the
sglang Gemma4Detector parsing algorithm (gemma_parser.py)
valid_name — every called tool actually exists
schema_valid — parsed arguments validate against the tool's JSON schema
stops_at_boundary — generation did not run past <tool_call|> into the
engine-owned <|tool_response> territory
thought_ok — (thinking scenarios) <|channel>thought opened and closed
before the first tool call
Usage:
python test_gemma4_scenarios.py [--scenario a|b|c|all] [--samples 4]
[--max-new-tokens 384] [--temperature 0.9]
[--model google/gemma-4-E2B-it]
[--report report.md]
"""
import argparse
import json
import time
import jsonschema
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import xgrammar as xgr
from gemma_parser import extract_tool_calls
DEFAULT_MODEL_ID = "google/gemma-4-E2B-it"
# ---------------------------------------------------------------- scenario A
CALENDAR_TOOL = {
"type": "function",
"function": {
"name": "create_calendar_event",
"description": "Create a calendar event and optionally send invites.",
"parameters": {
"type": "object",
"properties": {
"title": {"type": "string"},
"start": {"type": "string", "description": "ISO 8601 datetime"},
"duration_minutes": {"type": "integer"},
"priority": {"type": "string", "enum": ["low", "medium", "high"]},
"attendees": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": {"type": "string"},
"email": {"type": "string"},
},
"required": ["name", "email"],
},
},
"location": {
"type": "object",
"properties": {
"room": {"type": "string"},
"floor": {"type": "integer"},
},
"required": ["room"],
},
"send_invites": {"type": "boolean"},
},
"required": ["title", "start", "duration_minutes", "priority", "attendees"],
},
},
}
SCENARIO_A = {
"key": "a",
"name": "A. complex-schema",
"tools": [CALENDAR_TOOL],
"messages": [
{
"role": "user",
"content": (
'Schedule a high-priority meeting titled Q3 Roadmap Review ("final" draft) '
"on 2026-07-10 at 14:00 KST for 90 minutes in room Jupiter on the 7th floor. "
"Invite Alice Kim <alice.kim@example.com>, Bob Lee <bob.lee@example.com> and "
"Chris Park <chris.park@example.com>, and send the invites."
),
}
],
"reasoning": False,
"check_thought": False,
}
# ---------------------------------------------------------------- scenario B
WEATHER_TOOL = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a city.",
"parameters": {
"type": "object",
"properties": {
"city": {"type": "string"},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["city", "unit"],
},
},
}
CONVERT_TOOL = {
"type": "function",
"function": {
"name": "convert_temperature",
"description": "Convert a temperature value between units.",
"parameters": {
"type": "object",
"properties": {
"value": {"type": "number", "description": "numeric temperature value"},
"from_unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
"to_unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["value", "from_unit", "to_unit"],
},
},
}
SCENARIO_B = {
"key": "b",
"name": "B. multiturn-thinking",
"tools": [WEATHER_TOOL, CONVERT_TOOL],
"messages": [
{
"role": "user",
"content": (
"Compare the current weather in Seoul and Busan in celsius, one city at a "
"time. After both, also convert Seoul's temperature to fahrenheit with the "
"converter tool before answering."
),
},
{
"role": "assistant",
"tool_calls": [
{
"type": "function",
"function": {
"name": "get_weather",
"arguments": {"city": "Seoul", "unit": "celsius"},
},
}
],
},
{
"role": "tool",
"name": "get_weather",
"content": '{"temp_c": 3, "condition": "sunny", "humidity": 41}',
},
{
"role": "assistant",
"tool_calls": [
{
"type": "function",
"function": {
"name": "get_weather",
"arguments": {"city": "Busan", "unit": "celsius"},
},
}
],
},
{
"role": "tool",
"name": "get_weather",
"content": '{"temp_c": 8, "condition": "cloudy", "humidity": 63}',
},
],
"reasoning": True,
"check_thought": True,
}
# ---------------------------------------------------------------- scenario C
TICKET_TOOL = {
"type": "function",
"function": {
"name": "create_ticket",
"description": "Create a bug ticket in the issue tracker.",
"parameters": {
"type": "object",
"properties": {
"title": {"type": "string"},
"body": {
"type": "string",
"description": "Full description. Include configs/logs verbatim.",
},
"severity": {"type": "string", "enum": ["low", "medium", "high", "blocker"]},
"labels": {
"type": "array",
"items": {
"type": "string",
"enum": ["bug", "regression", "ui", "backend", "perf"],
},
},
"affected": {
"type": "object",
"properties": {
"service": {"type": "string"},
"version": {"type": "string"},
"regions": {"type": "array", "items": {"type": "string"}},
},
"required": ["service", "version"],
},
"cc_emails": {"type": "array", "items": {"type": "string"}},
"urgent_escalation": {"type": "boolean"},
"estimated_minutes": {"type": "integer"},
},
"required": ["title", "body", "severity", "labels", "affected"],
},
},
}
SEARCH_TICKETS_TOOL = {
"type": "function",
"function": {
"name": "search_tickets",
"description": "Search existing tickets.",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}, "limit": {"type": "integer"}},
"required": ["query"],
},
},
}
ESCALATE_TOOL = {
"type": "function",
"function": {
"name": "escalate_incident",
"description": "Escalate an existing incident by id.",
"parameters": {
"type": "object",
"properties": {"incident_id": {"type": "string"}, "level": {"type": "integer"}},
"required": ["incident_id", "level"],
},
},
}
SCENARIO_C = {
"key": "c",
"name": "C. adversarial-payload",
"tools": [TICKET_TOOL, SEARCH_TICKETS_TOOL, ESCALATE_TOOL],
"messages": [
{
"role": "user",
"content": (
"Our payments service is broken, this is critical! File a ticket titled "
"Payments retry storm after config rollout. The body must include, verbatim, "
'the deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], '
'"jitter": true}} and the error line: TypeError: cannot destructure '
"{id: undefined} at applyRetry (retry.js:42). It's a regression bug in the "
"backend, affecting service payments version 2.14.1 in regions ap-northeast-2 "
"and us-east-1. CC dev-alerts@example.com, escalate urgently, and estimate "
"45 minutes of work."
),
}
],
# Thinking is exercised in scenario B; here the budget goes to the payload itself.
# Elevated temperature with the full sampling tail (top_p=1.0) exposes format
# instability in the long string argument — settings a serving engine must survive.
"reasoning": False,
"check_thought": False,
"max_new_tokens": 640,
"temperature": 1.5,
"top_p": 1.0,
}
SCENARIOS = {"a": SCENARIO_A, "b": SCENARIO_B, "c": SCENARIO_C}
THOUGHT_BEGIN = "<|channel>thought"
THOUGHT_END = "<channel|>"
TOOL_CALL_START = "<|tool_call>"
TOOL_RESPONSE_START = "<|tool_response>"
METRIC_DESCRIPTIONS = {
"well_formed": "every tool call block is complete and parseable (sglang parser port)",
"valid_name": "every called tool exists",
"schema_valid": "parsed arguments validate against the tool's JSON schema",
"stops_at_boundary": (
"generation did not run past <tool_call|> into <|tool_response>. Per the Gemma 4 "
"prompt-formatting spec, the model is only responsible for the call up to "
"<tool_call|>; <|tool_response> is appended by the application with the real "
"tool result, and is registered as an additional stop sequence purely as a "
"backstop. This harness configures no such stop, so unconstrained runs "
"free-run past the boundary into engine-owned territory, while required-mode "
"grammar ends the call cleanly at an accept state."
),
"thought_ok": "<|channel>thought opened and closed before the first tool call",
}
def score_output(text: str, scenario: dict) -> dict:
tools_by_name = {t["function"]["name"]: t["function"]["parameters"] for t in scenario["tools"]}
# Anything past the <|tool_response> boundary is a separate problem (stops_at_boundary);
# parse only the part the serving engine would treat as the model's calls.
model_part = text.split(TOOL_RESPONSE_START)[0]
calls = extract_tool_calls(model_part)
complete = [(n, a) for n, a in calls if n is not None and a is not None]
well_formed = len(complete) > 0 and len(complete) == len(calls)
valid_name = well_formed and all(n in tools_by_name for n, _ in complete)
schema_valid = valid_name
error = None
if valid_name:
for name, args in complete:
try:
jsonschema.validate(args, tools_by_name[name])
except jsonschema.ValidationError as e:
schema_valid = False
error = f"{name}: {e.message}"
break
elif well_formed:
error = "unknown tool: " + ", ".join(n for n, _ in complete if n not in tools_by_name)
else:
error = "malformed or missing tool call block"
result = {
"well_formed": well_formed,
"valid_name": valid_name,
"schema_valid": schema_valid,
"stops_at_boundary": TOOL_RESPONSE_START not in text,
"calls": complete,
"error": error if not (well_formed and valid_name and schema_valid) else None,
}
if scenario["check_thought"]:
first_call = model_part.find(TOOL_CALL_START)
opened = model_part.startswith(THOUGHT_BEGIN)
closed = opened and THOUGHT_END in model_part[: first_call if first_call != -1 else None]
result["thought_ok"] = opened and closed
return result
def metric_keys_for(scenario: dict) -> list:
keys = ["well_formed", "valid_name", "schema_valid", "stops_at_boundary"]
if scenario["check_thought"]:
keys.append("thought_ok")
return keys
def run_scenario(scenario, model, tokenizer, args):
print(f"\n{'=' * 70}\n{scenario['name']} (reasoning={scenario['reasoning']})\n{'=' * 70}")
prompt = tokenizer.apply_chat_template(
scenario["messages"],
tools=scenario["tools"],
add_generation_prompt=True,
tokenize=False,
)
print("--- prompt tail ---")
print(prompt[-400:])
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
vocab_size = model.config.get_text_config().vocab_size
tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
compiler = xgr.GrammarCompiler(tokenizer_info)
stag = xgr.get_model_structural_tag(
model="gemma_4",
tools=scenario["tools"],
tool_choice="required",
reasoning=scenario["reasoning"],
)
compiled = compiler.compile_structural_tag(stag)
metric_keys = metric_keys_for(scenario)
max_new_tokens = max(args.max_new_tokens, scenario.get("max_new_tokens", 0))
temperature = scenario.get("temperature", args.temperature)
top_p = scenario.get("top_p", 0.95)
summary = {}
failures = []
for mode in ["constrained", "unconstrained"]:
counts = {k: 0 for k in metric_keys}
print(
f"\n--- {mode}: {args.samples} samples,"
f" temperature={temperature}, top_p={top_p} ---"
)
for s in range(args.samples):
torch.manual_seed(1000 + s)
gen_kwargs = dict(
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
)
if mode == "constrained":
gen_kwargs["logits_processor"] = [xgr.contrib.hf.LogitsProcessor(compiled)]
t0 = time.monotonic()
out = model.generate(**inputs, **gen_kwargs)
gen = tokenizer.decode(
out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False
)
gen = gen.split("<eos>")[0].split("<end_of_turn>")[0]
score = score_output(gen, scenario)
for k in metric_keys:
counts[k] += bool(score[k])
failed_keys = [k for k in metric_keys if not score[k]]
if failed_keys:
reasons = []
if score["error"]:
reasons.append(score["error"])
if "stops_at_boundary" in failed_keys:
reasons.append("ran past <tool_call|> into <|tool_response> after the call")
if "thought_ok" in failed_keys:
reasons.append("skipped or never closed the <|channel>thought section")
failures.append(
{
"mode": mode,
"sample": s,
"failed": failed_keys,
"reason": "; ".join(reasons),
"output": gen,
}
)
flags = " ".join(f"{k}={'O' if score[k] else 'X'}" for k in metric_keys)
print(f"\n[{mode} #{s}] ({time.monotonic() - t0:.0f}s) {flags}")
if score["error"]:
print(f" error: {score['error']}")
print(f" output: {gen[:700]!r}")
summary[mode] = counts
print(f"\n--- {scenario['name']} summary (out of {args.samples}) ---")
print(f"{'metric':<14}" + "".join(f"{m:>16}" for m in summary))
for k in metric_keys:
print(f"{k:<14}" + "".join(f"{summary[m][k]:>16}" for m in summary))
return {
"summary": summary,
"metric_keys": metric_keys,
"failures": failures,
"temperature": temperature,
"top_p": top_p,
"max_new_tokens": max_new_tokens,
"reasoning": scenario["reasoning"],
}
def write_report(path, model_id, args, results):
lines = [
"# Gemma-4 tool calling: constrained (xgrammar) vs unconstrained",
"",
f"- model: `{model_id}` via HF transformers",
'- constraint: `xgr.get_model_structural_tag("gemma_4", tools=..., '
'tool_choice="required")` compiled and applied with '
"`xgr.contrib.hf.LogitsProcessor`",
f"- {args.samples} samples per scenario per mode; identical sampling settings"
" for both modes (per-scenario values below)",
"- outputs validated with a pure-Python port of sglang's `Gemma4Detector`"
" parser plus `jsonschema`",
"",
"## Results (passing samples / total)",
"",
]
for name, res in results.items():
summary, metric_keys = res["summary"], res["metric_keys"]
lines.append(f"### {name}")
lines.append("")
lines.append(
f"temperature={res['temperature']}, top_p={res['top_p']},"
f" max_new_tokens={res['max_new_tokens']}, reasoning={res['reasoning']}"
)
lines.append("")
lines.append("| metric | constrained | unconstrained |")
lines.append("|---|---|---|")
for k in metric_keys:
c, u = summary["constrained"][k], summary["unconstrained"][k]
c_mark = "✅" if c == args.samples else "❌"
u_mark = "✅" if u == args.samples else "❌"
lines.append(f"| {k} | {c}/{args.samples} {c_mark} | {u}/{args.samples} {u_mark} |")
lines.append("")
lines += ["## Metric definitions", ""]
for k, desc in METRIC_DESCRIPTIONS.items():
lines.append(f"- `{k}` — {desc}")
lines.append("")
all_failures = [(name, f) for name, res in results.items() for f in res["failures"]]
if all_failures:
lines += ["## Failure examples", ""]
for name, f in all_failures[:12]:
failed = ", ".join(f"`{k}`" for k in f["failed"])
lines.append(f"**{name} / {f['mode']} #{f['sample']}** — failed {failed}: {f['reason']}")
lines.append("")
lines.append("```")
lines.append(f["output"][:900])
lines.append("```")
lines.append("")
with open(path, "w") as fp:
fp.write("\n".join(lines))
print(f"\nreport written to {path}")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--scenario", choices=["a", "b", "c", "all"], default="all")
parser.add_argument("--samples", type=int, default=4)
parser.add_argument("--max-new-tokens", type=int, default=384)
parser.add_argument("--temperature", type=float, default=0.9)
parser.add_argument("--model", default=DEFAULT_MODEL_ID)
parser.add_argument("--report", default=None, help="write a markdown report to this path")
args = parser.parse_args()
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"loading {args.model} on {device} ...")
tokenizer = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(
args.model, dtype=torch.bfloat16, device_map=device
)
model.eval()
selected = ["a", "b", "c"] if args.scenario == "all" else [args.scenario]
results = {}
for key in selected:
results[SCENARIOS[key]["name"]] = run_scenario(SCENARIOS[key], model, tokenizer, args)
print(f"\n{'=' * 70}\nOVERALL\n{'=' * 70}")
print(json.dumps({n: r["summary"] for n, r in results.items()}, indent=2, ensure_ascii=False))
if args.report:
write_report(args.report, args.model, args, results)
if __name__ == "__main__":
main()
test_gemma4_xgrammar.py
+120
View File
@@ -0,0 +1,120 @@
"""End-to-end test: constrained decoding for Gemma-4 tool calling with the modified xgrammar.
Loads google/gemma-4-E2B-it via HF transformers, builds the gemma_4 builtin
structural tag from xgrammar, and generates with xgrammar's LogitsProcessor.
Verifies:
1. The grammar compiles against the real Gemma-4 tokenizer.
2. Constrained generation emits a well-formed <|tool_call>call:name{...}<tool_call|>.
3. The emitted arguments use Gemma's <|"|> string delimiters and satisfy the schema.
4. (Alignment) Unconstrained generation is compared against the same grammar.
"""
import json
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import xgrammar as xgr
from xgrammar.testing import _is_grammar_accept_string
MODEL_ID = "google/gemma-4-E2B-it"
TOOLS = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a city.",
"parameters": {
"type": "object",
"properties": {
"city": {"type": "string", "description": "City name"},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["city"],
},
},
}
]
MESSAGES = [{"role": "user", "content": "What's the weather in Seoul in celsius?"}]
def main():
print(f"=== loading tokenizer/model: {MODEL_ID} ===")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, device_map="mps")
model.eval()
vocab_size = model.config.get_text_config().vocab_size
print(f"vocab_size={vocab_size}")
prompt = tokenizer.apply_chat_template(
MESSAGES, tools=TOOLS, add_generation_prompt=True, tokenize=False
)
print("=== rendered prompt (tail) ===")
print(prompt[-600:])
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("mps")
tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
compiler = xgr.GrammarCompiler(tokenizer_info)
results = {}
for tool_choice, reasoning in [("required", False), ("auto", True)]:
label = f"tool_choice={tool_choice}, reasoning={reasoning}"
print(f"\n=== structural tag: {label} ===")
stag = xgr.get_model_structural_tag(
model="gemma_4", tools=TOOLS, tool_choice=tool_choice, reasoning=reasoning
)
t0 = time.monotonic()
compiled = compiler.compile_structural_tag(stag)
print(f"grammar compile time: {time.monotonic() - t0:.2f}s")
processor = xgr.contrib.hf.LogitsProcessor(compiled)
t0 = time.monotonic()
out = model.generate(
**inputs, max_new_tokens=128, do_sample=False, logits_processor=[processor]
)
gen = tokenizer.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False)
print(f"generate time: {time.monotonic() - t0:.1f}s")
print(f"--- constrained output ({label}) ---")
print(repr(gen))
results[label] = gen
print("\n=== unconstrained (alignment check) ===")
out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
unconstrained = tokenizer.decode(
out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False
)
print(repr(unconstrained))
results["unconstrained"] = unconstrained
print("\n=== verification ===")
ok = True
required_out = results["tool_choice=required, reasoning=False"]
# strip trailing special tokens after the tool call for the checks
body = required_out.split("<tool_call|>")[0] + "<tool_call|>" if "<tool_call|>" in required_out else required_out
checks = [
("required: contains <|tool_call>call:get_weather{", "<|tool_call>call:get_weather{" in body),
("required: closes with }<tool_call|>", body.rstrip().endswith("<tool_call|>")),
('required: city uses <|"|> delimiters', '<|"|>' in body),
("required: no JSON-quoted args", '"city"' not in body),
]
stag_req = xgr.get_model_structural_tag(
model="gemma_4", tools=TOOLS, tool_choice="required", reasoning=False
)
g = xgr.Grammar.from_structural_tag(stag_req)
checks.append(("required: output accepted by grammar", _is_grammar_accept_string(g, body)))
for name, passed in checks:
print(f" [{'PASS' if passed else 'FAIL'}] {name}")
ok &= passed
print("\nRESULT:", "ALL PASS" if ok else "SOME CHECKS FAILED")
print("\nJSON summary:")
print(json.dumps({k: v for k, v in results.items()}, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()