From 9484167719e3af2dbdb58a5369f6c5ba2b44c380 Mon Sep 17 00:00:00 2001 From: "wanhae.lee" Date: Sun, 5 Jul 2026 00:33:50 +0900 Subject: [PATCH] init --- .gitignore | 14 + Makefile | 37 +++ README.md | 65 +++++ gemma_parser.py | 215 +++++++++++++++ report.md | 128 +++++++++ test_gemma4_scenarios.py | 546 +++++++++++++++++++++++++++++++++++++++ test_gemma4_xgrammar.py | 120 +++++++++ 7 files changed, 1125 insertions(+) create mode 100644 .gitignore create mode 100644 Makefile create mode 100644 README.md create mode 100644 gemma_parser.py create mode 100644 report.md create mode 100644 test_gemma4_scenarios.py create mode 100644 test_gemma4_xgrammar.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..00179c6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,14 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ + +# Virtual environments +.venv/ +venv/ + +# uv +uv.lock + +# OS +.DS_Store diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..330bcad --- /dev/null +++ b/Makefile @@ -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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..f943bf6 --- /dev/null +++ b/README.md @@ -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{...}` 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 `` 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_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` + 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` diff --git a/gemma_parser.py b/gemma_parser.py new file mode 100644 index 0000000..46c39c1 --- /dev/null +++ b/gemma_parser.py @@ -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 = "" +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 diff --git a/report.md b/report.md new file mode 100644 index 0000000..6c9b057 --- /dev/null +++ b/report.md @@ -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 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 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_response> +``` + +**A. complex-schema / unconstrained #1** — failed `stops_at_boundary`: ran past 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_response> +``` + +**A. complex-schema / unconstrained #2** — failed `stops_at_boundary`: ran past 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_response> +``` + +**A. complex-schema / unconstrained #3** — failed `stops_at_boundary`: ran past 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_response> +``` + +**B. multiturn-thinking / unconstrained #0** — failed `stops_at_boundary`, `thought_ok`: ran past 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_response> +``` + +**B. multiturn-thinking / unconstrained #1** — failed `stops_at_boundary`, `thought_ok`: ran past 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_response> +``` + +**B. multiturn-thinking / unconstrained #2** — failed `stops_at_boundary`, `thought_ok`: ran past 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_response> +``` + +**B. multiturn-thinking / unconstrained #3** — failed `stops_at_boundary`, `thought_ok`: ran past 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_response> +``` + +**C. adversarial-payload / unconstrained #0** — failed `stops_at_boundary`: ran past 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_response> +``` + +**C. adversarial-payload / unconstrained #1** — failed `stops_at_boundary`: ran past 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_response> +``` + +**C. adversarial-payload / unconstrained #2** — failed `stops_at_boundary`: ran past 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_response> +``` + +**C. adversarial-payload / unconstrained #3** — failed `stops_at_boundary`: ran past 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_response> +``` diff --git a/test_gemma4_scenarios.py b/test_gemma4_scenarios.py new file mode 100644 index 0000000..48ad496 --- /dev/null +++ b/test_gemma4_scenarios.py @@ -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 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 , Bob Lee and " + "Chris Park , 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 = "" +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 into <|tool_response>. Per the Gemma 4 " + "prompt-formatting spec, the model is only responsible for the call up to " + "; <|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("")[0].split("")[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 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() diff --git a/test_gemma4_xgrammar.py b/test_gemma4_xgrammar.py new file mode 100644 index 0000000..bafb1f8 --- /dev/null +++ b/test_gemma4_xgrammar.py @@ -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{...}. + 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("")[0] + "" if "" in required_out else required_out + checks = [ + ("required: contains <|tool_call>call:get_weather{", "<|tool_call>call:get_weather{" in body), + ("required: closes with }", body.rstrip().endswith("")), + ('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()