init

2026-07-05 00:33:50 +09:00
commit 9484167719
7 changed files with 1125 additions and 0 deletions
@@ -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{...}<tool_call|>.
+  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("<tool_call|>")[0] + "<tool_call|>" if "<tool_call|>" in required_out else required_out
+    checks = [
+        ("required: contains <|tool_call>call:get_weather{", "<|tool_call>call:get_weather{" in body),
+        ("required: closes with }<tool_call|>", body.rstrip().endswith("<tool_call|>")),
+        ('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()