"""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 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 /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 , Bob Lee and " "Chris Park , 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 = "" 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=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("")[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]] 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") 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("")[0].split("")[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 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()