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