wan/xgrammar-gemma4-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
xgrammar-gemma4-test/test_gemma4_scenarios.py
T
wanhae.lee 42c3db9cd4 ...
2026-07-05 03:02:25 +09:00

820 lines
33 KiB
Python
Raw Permalink Blame History

"""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 with the full
sampling tail (top_p=1.0) to expose format instability.
D. structured-extraction — a real agentic loop (not a single canned
generation): for each of 3 raw Wikipedia scrapes
loaded straight from test/ (LY Corporation history
from lineyahoo.txt, Z Intermediate Global infobox
from z_holdings.txt, LINE app infobox from
line.txt, Korean-language), the harness calls
model.generate() once per citation/footnote marker
the model extracts, feeding its own actual prior
tool_calls/tool-acks back as context, until it
signals finish_document and moves to the next doc.
Scores every individual turn (dozens of real
generate() calls per run) — a model's own
imperfect output compounding in its own context is
a more realistic stress test than one hand-written
complex payload.
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]
[--temperature 0.9]
[--model google/gemma-4-E2B-it]
[--report report.md] [--quantize]
"""
import argparse
import json
import os
import time
import jsonschema
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import xgrammar as xgr
from gemma_parser import extract_tool_calls
DEFAULT_MODEL_ID = "google/gemma-4-E2B-it"
TEST_DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test")
# No real cap intended — just high enough that HF's tiny default max_length (20) never
# truncates a call before it reaches <end_of_turn>/eos on its own.
MAX_NEW_TOKENS = 4096
def _load_infobox(filename: str, start_marker: str, end_marker: str) -> str:
"""Slice the infobox block out of a raw Wikipedia scrape in test/."""
path = os.path.join(TEST_DATA_DIR, filename)
with open(path, encoding="utf-8") as fp:
lines = fp.read().splitlines()
start = next(i for i, line in enumerate(lines) if line.startswith(start_marker))
end = next(i for i, line in enumerate(lines) if i > start and line.startswith(end_marker))
return "\n".join(lines[start : end + 1])
def _load_section(filename: str, start_line: str, end_line: str) -> str:
"""Slice an exact-match line range (exclusive of end_line) out of a scrape in test/."""
path = os.path.join(TEST_DATA_DIR, filename)
with open(path, encoding="utf-8") as fp:
lines = fp.read().splitlines()
start = next(i for i, line in enumerate(lines) if line.strip() == start_line)
end = next(i for i, line in enumerate(lines) if i > start and line.strip() == end_line)
return "\n".join(lines[start:end])
# ---------------------------------------------------------------- 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."
),
}
],
"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}',
},
],
"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."
),
}
],
# Reasoning is enabled for every scenario; temperature always comes from the
# Makefile-configured value (see run_scenario). The full sampling tail (top_p=1.0)
# exposes format instability in the long string argument — settings a serving
# engine must survive.
"check_thought": False,
"top_p": 1.0,
}
# ---------------------------------------------------------------- scenario D
EXTRACT_CITATION_TOOL = {
"type": "function",
"function": {
"name": "extract_citation",
"description": (
"Record one citation/footnote marker found in the source text, verbatim, "
"along with a short note about what it's attached to."
),
"parameters": {
"type": "object",
"properties": {
"marker": {
"type": "string",
"description": "the citation marker exactly as written, e.g. [17] or [広報 3]",
},
"note": {
"type": "string",
"description": "short verbatim snippet of the fact/sentence the citation is attached to",
},
},
"required": ["marker", "note"],
},
},
}
FINISH_DOCUMENT_TOOL = {
"type": "function",
"function": {
"name": "finish_document",
"description": "Signal that every citation marker in the current document has been extracted.",
"parameters": {"type": "object", "properties": {}, "required": []},
},
}
LY_HISTORY = _load_section("lineyahoo.txt", "1990年代", "2010年代")
Z_INFOBOX = _load_infobox("z_holdings.txt", "種類", "テンプレートを表示")
LINE_APP_INFOBOX = _load_infobox("line.txt", "개발", "Chrome")
# (label, source filename, raw text, turn budget = citation count + slack)
SCENARIO_D_DOCS = [
("LY Corporation history timeline", "lineyahoo.txt", LY_HISTORY, 12),
("Z Intermediate Global infobox", "z_holdings.txt", Z_INFOBOX, 4),
("LINE app infobox (Korean-language scrape)", "line.txt", LINE_APP_INFOBOX, 3),
]
def _scenario_d_doc_prompt(label: str, source_file: str, text: str, is_first_doc: bool) -> str:
header = "Below is a raw Wikipedia scrape" if is_first_doc else "Now do the same for the next document"
return (
f"{header} ({label}, source: {source_file}). It contains citation/footnote "
"markers like [17] or [広報 3]. Go through it one marker at a time: each turn, "
"find the NEXT marker you haven't already extracted (check what's already in our "
"conversation) and call extract_citation with that exact marker and a short note "
"about what it's attached to — one marker per turn. Once every marker in this "
"document has been extracted, call finish_document instead.\n\n"
f"{text}"
)
SCENARIO_D = {
"key": "d",
"name": "D. structured-extraction",
"tools": [EXTRACT_CITATION_TOOL, FINISH_DOCUMENT_TOOL],
# A real agentic loop, not hardcoded context: for each of the 3 documents (all loaded
# straight from test/), we call model.generate() repeatedly — one real turn per
# citation marker — feeding the model's own actual prior tool_calls/tool-acks back as
# context, until it calls finish_document, then move to the next document. This scores
# every individual turn's wire-format adherence (dozens of real generate() calls per
# run) rather than one clean single-shot generation, since a model's own imperfect
# outputs compounding in its own context is a more realistic stress test than a single
# hand-written complex payload.
"docs": SCENARIO_D_DOCS,
"agentic": True,
"check_thought": False,
"top_p": 0.95,
}
SCENARIOS = {"a": SCENARIO_A, "b": SCENARIO_B, "c": SCENARIO_C, "d": SCENARIO_D}
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=True)\n{'=' * 70}")
prompt = tokenizer.apply_chat_template(
scenario["messages"],
tools=scenario["tools"],
add_generation_prompt=True,
tokenize=False,
enable_thinking=True,
)
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=True,
)
compiled = compiler.compile_structural_tag(stag)
metric_keys = metric_keys_for(scenario)
temperature = args.temperature
top_p = scenario.get("top_p", 0.95)
summary = {}
samples = []
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)]
print(f"\n[{mode} #{s}] streaming:")
gen_kwargs["streamer"] = TextStreamer(
tokenizer, skip_prompt=True, skip_special_tokens=False
)
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]]
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")
samples.append(
{
"mode": mode,
"sample": s,
"passed": not failed_keys,
"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,
"samples": samples,
"temperature": temperature,
"top_p": top_p,
"reasoning": True,
}
def run_agentic_scenario(scenario, model, tokenizer, args):
"""Real multi-turn agent loop: one model.generate() call per citation marker per
document, feeding the model's own actual outputs back as context, moving to the next
document once it calls finish_document. Scores every individual turn."""
print(f"\n{'=' * 70}\n{scenario['name']} (agentic)\n{'=' * 70}")
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=True
)
compiled = compiler.compile_structural_tag(stag)
metric_keys = ["well_formed", "valid_name", "schema_valid", "stops_at_boundary"]
temperature = args.temperature
top_p = scenario.get("top_p", 0.95)
summary = {}
samples = []
for mode in ["constrained", "unconstrained"]:
counts = {k: 0 for k in metric_keys}
total_turns = 0
print(f"\n--- {mode}: {args.samples} samples, temperature={temperature}, top_p={top_p} ---")
for s in range(args.samples):
torch.manual_seed(3000 + s)
messages = []
for doc_idx, (label, source_file, text, turn_budget) in enumerate(scenario["docs"]):
messages.append(
{
"role": "user",
"content": _scenario_d_doc_prompt(label, source_file, text, doc_idx == 0),
}
)
for turn in range(turn_budget):
prompt = tokenizer.apply_chat_template(
messages,
tools=scenario["tools"],
add_generation_prompt=True,
tokenize=False,
enable_thinking=True,
)
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(
model.device
)
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)]
label_short = f"{s}/{label[:24]}/turn{turn}"
print(f"\n[{mode} {label_short}] streaming:")
gen_kwargs["streamer"] = TextStreamer(
tokenizer, skip_prompt=True, skip_special_tokens=False
)
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)
total_turns += 1
for k in metric_keys:
counts[k] += bool(score[k])
failed_keys = [k for k in metric_keys if not score[k]]
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")
samples.append(
{
"mode": mode,
"sample": label_short,
"passed": not failed_keys,
"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} {label_short}] ({time.monotonic() - t0:.0f}s) {flags}")
if score["error"]:
print(f" error: {score['error']}")
print(f" output: {gen[:300]!r}")
model_part = gen.split(TOOL_RESPONSE_START)[0]
parsed = extract_tool_calls(model_part)
complete = [(n, a) for n, a in parsed if n is not None and a is not None]
if not complete:
print(" (derailed: no parseable call, abandoning this document)")
break
messages.append(
{
"role": "assistant",
"tool_calls": [
{"type": "function", "function": {"name": n, "arguments": a}}
for n, a in complete
],
}
)
for n, a in complete:
messages.append({"role": "tool", "name": n, "content": '{"status": "ok"}'})
if any(n == "finish_document" for n, _ in complete):
break
summary[mode] = counts
summary[mode]["_total"] = total_turns
print(f"\n--- {scenario['name']} summary ---")
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]}/{summary[m]['_total']:>10}" for m in summary))
return {
"summary": summary,
"metric_keys": metric_keys,
"samples": samples,
"temperature": temperature,
"top_p": top_p,
"reasoning": True,
}
def write_report(path, model_id, args, results):
quant_note = "torchao int8 weight-only" if args.quantize else "bf16 full precision"
lines = [
"# Gemma-4 tool calling: constrained (xgrammar) vs unconstrained",
"",
f"- model: `{model_id}` via HF transformers ({quant_note})",
'- 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" reasoning={res['reasoning']}"
)
lines.append("")
lines.append("| metric | constrained | unconstrained |")
lines.append("|---|---|---|")
c_total = summary["constrained"].get("_total", args.samples)
u_total = summary["unconstrained"].get("_total", args.samples)
for k in metric_keys:
c, u = summary["constrained"][k], summary["unconstrained"][k]
c_mark = "✅" if c == c_total else "❌"
u_mark = "✅" if u == u_total else "❌"
lines.append(f"| {k} | {c}/{c_total} {c_mark} | {u}/{u_total} {u_mark} |")
lines.append("")
lines += ["## Metric definitions", ""]
for k, desc in METRIC_DESCRIPTIONS.items():
lines.append(f"- `{k}` — {desc}")
lines.append("")
lines += ["## Sample outputs (raw)", ""]
for name, res in results.items():
lines.append(f"### {name}")
lines.append("")
for s in res["samples"]:
status = "PASS ✅" if s["passed"] else "FAIL ❌"
header = f"**{name} / {s['mode']} #{s['sample']}** — {status}"
if not s["passed"]:
failed = ", ".join(f"`{k}`" for k in s["failed"])
header += f" (failed {failed}: {s['reason']})"
lines.append(header)
lines.append("")
lines.append("```")
lines.append(s["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", "d", "all"], default="all")
parser.add_argument("--samples", type=int, default=4)
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")
parser.add_argument(
"--quantize",
action="store_true",
help="load the model with torchao int8 weight-only quantization (opt-in; default is bf16)",
)
args = parser.parse_args()
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"loading {args.model} on {device} (quantize={args.quantize}) ...")
tokenizer = AutoTokenizer.from_pretrained(args.model)
if args.quantize:
from torchao.quantization import Int8WeightOnlyConfig
from transformers import TorchAoConfig
quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
model = AutoModelForCausalLM.from_pretrained(
args.model,
dtype=torch.bfloat16,
device_map=device,
quantization_config=quantization_config,
)
else:
model = AutoModelForCausalLM.from_pretrained(
args.model, dtype=torch.bfloat16, device_map=device
)
model.eval()
selected = ["a", "b", "c", "d"] if args.scenario == "all" else [args.scenario]
results = {}
for key in selected:
scenario = SCENARIOS[key]
runner = run_agentic_scenario if scenario.get("agentic") else run_scenario
results[scenario["name"]] = runner(scenario, 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:
report_path = args.report
if args.quantize:
root, ext = os.path.splitext(report_path)
report_path = f"{root}.int8{ext}"
write_report(report_path, args.model, args, results)
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink