wanhae.lee
547 lines
21 KiB
Python
547 lines
21 KiB
Python
"""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()
|