...
This commit is contained in:
wanhae.lee
+303
-35
@@ -14,8 +14,23 @@ google/gemma-4-E2B-it, across scenarios designed to break tool call syntax:
|
||||
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.
|
||||
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
|
||||
@@ -29,23 +44,48 @@ Metrics per sample:
|
||||
|
||||
Usage:
|
||||
python test_gemma4_scenarios.py [--scenario a|b|c|all] [--samples 4]
|
||||
[--max-new-tokens 384] [--temperature 0.9]
|
||||
[--temperature 0.9]
|
||||
[--model google/gemma-4-E2B-it]
|
||||
[--report report.md]
|
||||
[--report report.md] [--quantize]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import jsonschema
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
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
|
||||
|
||||
@@ -102,7 +142,6 @@ SCENARIO_A = {
|
||||
),
|
||||
}
|
||||
],
|
||||
"reasoning": False,
|
||||
"check_thought": False,
|
||||
}
|
||||
|
||||
@@ -189,7 +228,6 @@ SCENARIO_B = {
|
||||
"content": '{"temp_c": 8, "condition": "cloudy", "humidity": 63}',
|
||||
},
|
||||
],
|
||||
"reasoning": True,
|
||||
"check_thought": True,
|
||||
}
|
||||
|
||||
@@ -279,17 +317,94 @@ SCENARIO_C = {
|
||||
),
|
||||
}
|
||||
],
|
||||
# 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,
|
||||
# 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,
|
||||
"max_new_tokens": 640,
|
||||
"temperature": 1.5,
|
||||
"top_p": 1.0,
|
||||
}
|
||||
|
||||
SCENARIOS = {"a": SCENARIO_A, "b": SCENARIO_B, "c": SCENARIO_C}
|
||||
# ---------------------------------------------------------------- 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|>"
|
||||
@@ -363,13 +478,14 @@ def metric_keys_for(scenario: dict) -> list:
|
||||
|
||||
|
||||
def run_scenario(scenario, model, tokenizer, args):
|
||||
print(f"\n{'=' * 70}\n{scenario['name']} (reasoning={scenario['reasoning']})\n{'=' * 70}")
|
||||
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:])
|
||||
@@ -382,13 +498,12 @@ def run_scenario(scenario, model, tokenizer, args):
|
||||
model="gemma_4",
|
||||
tools=scenario["tools"],
|
||||
tool_choice="required",
|
||||
reasoning=scenario["reasoning"],
|
||||
reasoning=True,
|
||||
)
|
||||
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)
|
||||
temperature = args.temperature
|
||||
top_p = scenario.get("top_p", 0.95)
|
||||
summary = {}
|
||||
samples = []
|
||||
@@ -401,13 +516,17 @@ def run_scenario(scenario, model, tokenizer, args):
|
||||
for s in range(args.samples):
|
||||
torch.manual_seed(1000 + s)
|
||||
gen_kwargs = dict(
|
||||
max_new_tokens=max_new_tokens,
|
||||
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(
|
||||
@@ -452,16 +571,141 @@ def run_scenario(scenario, model, tokenizer, args):
|
||||
"samples": samples,
|
||||
"temperature": temperature,
|
||||
"top_p": top_p,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"reasoning": scenario["reasoning"],
|
||||
"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",
|
||||
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`",
|
||||
@@ -479,16 +723,18 @@ def write_report(path, model_id, args, results):
|
||||
lines.append("")
|
||||
lines.append(
|
||||
f"temperature={res['temperature']}, top_p={res['top_p']},"
|
||||
f" max_new_tokens={res['max_new_tokens']}, reasoning={res['reasoning']}"
|
||||
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 == args.samples else "❌"
|
||||
u_mark = "✅" if u == args.samples else "❌"
|
||||
lines.append(f"| {k} | {c}/{args.samples} {c_mark} | {u}/{args.samples} {u_mark} |")
|
||||
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", ""]
|
||||
@@ -520,31 +766,53 @@ def write_report(path, model_id, args, results):
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--scenario", choices=["a", "b", "c", "all"], default="all")
|
||||
parser.add_argument("--scenario", choices=["a", "b", "c", "d", "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")
|
||||
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} ...")
|
||||
print(f"loading {args.model} on {device} (quantize={args.quantize}) ...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model, dtype=torch.bfloat16, device_map=device
|
||||
)
|
||||
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"] if args.scenario == "all" else [args.scenario]
|
||||
selected = ["a", "b", "c", "d"] if args.scenario == "all" else [args.scenario]
|
||||
results = {}
|
||||
for key in selected:
|
||||
results[SCENARIOS[key]["name"]] = run_scenario(SCENARIOS[key], model, tokenizer, args)
|
||||
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:
|
||||
write_report(args.report, args.model, args, results)
|
||||
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__":
|
||||
|
||||
Reference in New Issue
Block a user