...

test_gemma4_xgrammar.py

@@ -10,6 +10,7 @@ Verifies:
 """
 
 import json
+import os
 import time
 
 import torch
@@ -19,6 +20,7 @@ import xgrammar as xgr
 from xgrammar.testing import _is_grammar_accept_string
 
 MODEL_ID = "google/gemma-4-E2B-it"
+QUANTIZE_INT8 = os.environ.get("QUANTIZE_INT8", "0") == "1"
 
 TOOLS = [
     {
@@ -44,14 +46,31 @@ 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")
+    if QUANTIZE_INT8:
+        from torchao.quantization import Int8WeightOnlyConfig
+        from transformers import TorchAoConfig
+
+        print("=== loading with torchao int8 weight-only quantization ===")
+        quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            dtype=torch.bfloat16,
+            device_map="mps",
+            quantization_config=quantization_config,
+        )
+    else:
+        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}")
 
+    # reasoning=False prompt/inputs; also used below for the unconstrained alignment
+    # check, which is compared against the tool_choice=required, reasoning=False case.
     prompt = tokenizer.apply_chat_template(
-        MESSAGES, tools=TOOLS, add_generation_prompt=True, tokenize=False
+        MESSAGES, tools=TOOLS, add_generation_prompt=True, tokenize=False, enable_thinking=False
     )
     print("=== rendered prompt (tail) ===")
     print(prompt[-600:])
@@ -71,12 +90,27 @@ def main():
         compiled = compiler.compile_structural_tag(stag)
         print(f"grammar compile time: {time.monotonic() - t0:.2f}s")
 
+        # Rebuild the prompt so <|think|> is present exactly when this iteration's
+        # grammar expects a thought block (enable_thinking must match `reasoning`).
+        iter_prompt = tokenizer.apply_chat_template(
+            MESSAGES,
+            tools=TOOLS,
+            add_generation_prompt=True,
+            tokenize=False,
+            enable_thinking=reasoning,
+        )
+        iter_inputs = tokenizer(iter_prompt, return_tensors="pt", add_special_tokens=False).to(
+            "mps"
+        )
+
         processor = xgr.contrib.hf.LogitsProcessor(compiled)
         t0 = time.monotonic()
         out = model.generate(
-            **inputs, max_new_tokens=128, do_sample=False, logits_processor=[processor]
+            **iter_inputs, max_new_tokens=128, do_sample=False, logits_processor=[processor]
+        )
+        gen = tokenizer.decode(
+            out[0][iter_inputs["input_ids"].shape[1] :], skip_special_tokens=False
         )
-        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))