| """Same-model baseline: the base Qwen3.5-0.8B (the instruct/chat release deg starts from, |
| NOT deg's weights) writes the whole JSON object under xgrammar JSON-schema-constrained |
| greedy decoding (HF transformers generate + xgrammar logits processor), batch 1. |
| |
| Prompt: chat template, thinking off, :data:`deg.bench.systems.INSTRUCTION` as the system |
| message, schema + context as the user message. The grammar is compiled from the schema |
| (``x-generate`` dropped as an annotation) with standard ``", "`` / ``": "`` separators and no |
| other whitespace; ``compact=True`` uses ``","`` / ``":"`` (fewer tokens, but the 0.8B model |
| then often writes a lone ``","`` as a string value: an ablation, not the main arm). Grammar compilation is timed separately (``extra.compile_s``) and is not in |
| ``latency_s``; generated tokens count every emitted token except the stop token. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
|
|
| from deg.bench.score import Prediction |
| from deg.bench.systems import INSTRUCTION, System, user_message |
|
|
| MAX_NEW_TOKENS = 1536 |
|
|
|
|
| def grammar_schema(schema: Any, drop: tuple[str, ...] = ("x-generate",)) -> Any: |
| def walk(node: Any) -> Any: |
| if isinstance(node, dict): |
| return {k: walk(v) for k, v in node.items() if k not in drop} |
| if isinstance(node, list): |
| return [walk(v) for v in node] |
| return node |
| return walk(copy.deepcopy(schema)) |
|
|
|
|
| def stop_ids(tok: Any, model_dir: Path) -> list[int]: |
| """End-of-turn ids: Qwen's ``<|im_end|>`` when the vocab has it, the tokenizer eos, and |
| every ``eos_token_id`` in ``generation_config.json`` (Gemma / LFM / MiniCPM turn ends).""" |
| stop: list[int] = [] |
| im_end = tok.convert_tokens_to_ids("<|im_end|>") |
| if isinstance(im_end, int) and im_end != tok.unk_token_id: |
| stop.append(im_end) |
| cands: list[Any] = [tok.eos_token_id] |
| gc = model_dir / "generation_config.json" |
| if gc.exists(): |
| e = json.loads(gc.read_text()).get("eos_token_id") |
| cands += e if isinstance(e, list) else [e] |
| for t in cands: |
| if isinstance(t, int) and t not in stop: |
| stop.append(t) |
| return stop |
|
|
|
|
| def native_context(model_dir: Path) -> int | None: |
| """The model's own maximum context (``max_position_embeddings``, text config first).""" |
| cfg = json.loads((model_dir / "config.json").read_text()) |
| for c in (cfg.get("text_config") or {}, cfg): |
| if c.get("max_position_embeddings"): |
| return int(c["max_position_embeddings"]) |
| return None |
|
|
|
|
| def chat_prompt(tok: Any, rec: dict[str, Any]) -> str: |
| """The model's own chat template, thinking off where the template supports it; a |
| template without a system role gets the instruction prepended to the user turn.""" |
| user = user_message(rec["context"], rec["schema"]) |
| try: |
| return tok.apply_chat_template( |
| [{"role": "system", "content": INSTRUCTION}, {"role": "user", "content": user}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| except Exception: |
| return tok.apply_chat_template( |
| [{"role": "user", "content": INSTRUCTION + "\n\n" + user}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
|
|
|
|
| class ConstrainedSystem(System): |
| def __init__(self, model_path: str | None, label: str | None = None, |
| compact: bool = False) -> None: |
| import xgrammar as xgr |
| from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen3_5ForCausalLM |
|
|
| from deg.model import base_snapshot |
|
|
| path = Path(model_path) if model_path else base_snapshot() |
| self.tok = AutoTokenizer.from_pretrained(str(path)) |
| model_type = json.loads((path / "config.json").read_text()).get("model_type", "") |
| cls = Qwen3_5ForCausalLM if model_type.startswith("qwen3_5") else AutoModelForCausalLM |
| self.lm = cls.from_pretrained(str(path), dtype=torch.bfloat16) |
| self.lm.to("cuda").eval() |
| vocab = self.lm.get_output_embeddings().weight.shape[0] |
| self.stop = stop_ids(self.tok, path) |
| self.native_ctx = native_context(path) |
| info = xgr.TokenizerInfo.from_huggingface(self.tok, vocab_size=vocab, |
| stop_token_ids=self.stop) |
| self.compiler = xgr.GrammarCompiler(info, max_threads=8) |
| self.xgr = xgr |
| self.separators = (",", ":") if compact else (", ", ": ") |
| self.name = label or ("qwen3.5-0.8b-xgrammar" + ("-compact" if compact else "")) |
| self.info = {"model": str(path), |
| "weights": ("base Qwen3.5-0.8B (chat release), not deg" |
| if model_path is None else f"{model_type} (as released)"), |
| "decoding": "greedy, xgrammar JSON-schema constrained", |
| "separators": list(self.separators), "native_context": self.native_ctx, |
| "dtype": "bf16", "batch": 1, "runtime": "HF transformers generate"} |
| self._warm() |
|
|
| def _warm(self) -> None: |
| rec = {"uid": "warm", "context": "Order AB-1 arrived late.", |
| "schema": {"type": "object", "properties": {"id": {"type": "string"}}, |
| "required": ["id"]}} |
| self.predict(rec) |
|
|
| def prompt_ids(self, rec: dict[str, Any]) -> list[int]: |
| return self.tok(chat_prompt(self.tok, rec), add_special_tokens=False)["input_ids"] |
|
|
| def compile(self, schema: Any) -> tuple[Any, str]: |
| clean = grammar_schema(schema) |
| try: |
| return self.compiler.compile_json_schema( |
| json.dumps(clean), any_whitespace=False, separators=self.separators), "full" |
| except Exception: |
| loose = grammar_schema(clean, drop=("format", "pattern")) |
| return self.compiler.compile_json_schema( |
| json.dumps(loose), any_whitespace=False, separators=self.separators), "no-format" |
|
|
| @torch.no_grad() |
| def predict(self, rec: dict[str, Any]) -> Prediction: |
| from xgrammar.contrib.hf import LogitsProcessor |
|
|
| t0 = time.perf_counter() |
| try: |
| grammar, grammar_mode = self.compile(rec["schema"]) |
| except Exception as exc: |
| return Prediction(rec["uid"], error=f"grammar: {str(exc)[:200]}") |
| compile_s = time.perf_counter() - t0 |
| torch.cuda.synchronize() |
| t1 = time.perf_counter() |
| ids = self.prompt_ids(rec) |
| if self.native_ctx and len(ids) >= self.native_ctx: |
| |
| return Prediction(rec["uid"], error=f"context_exceeded: native {len(ids)} prompt " |
| f"tokens >= {self.native_ctx}", |
| prompt_tokens=len(ids)) |
| inp = torch.tensor([ids], device="cuda") |
| out = self.lm.generate(inp, attention_mask=torch.ones_like(inp), |
| logits_processor=[LogitsProcessor(grammar)], |
| max_new_tokens=MAX_NEW_TOKENS, do_sample=False, |
| eos_token_id=self.stop, pad_token_id=self.stop[0]) |
| torch.cuda.synchronize() |
| dt = time.perf_counter() - t1 |
| new = out[0, len(ids):].tolist() |
| while new and new[-1] in self.stop: |
| new.pop() |
| text = self.tok.decode(new, skip_special_tokens=True) |
| extra = {"compile_s": round(compile_s, 4), "grammar": grammar_mode} |
| try: |
| value = json.loads(text) |
| except json.JSONDecodeError as exc: |
| hit_max = len(new) >= MAX_NEW_TOKENS - 1 |
| return Prediction(rec["uid"], error=f"json: {exc}"[:200] + (" (max tokens)" |
| if hit_max else ""), |
| generated_tokens=len(new), prompt_tokens=len(ids), latency_s=dt, |
| extra={**extra, "text": text[-300:]}) |
| return Prediction(rec["uid"], value=value, generated_tokens=len(new), |
| prompt_tokens=len(ids), latency_s=dt, extra=extra) |
|
|
|
|
| class LlamaCppConstrainedSystem(System): |
| """The same baseline on llama.cpp: a GGUF model on a llama.cpp server (CUDA), the |
| model's own chat template (from its HF snapshot, ``model_path``; default the base |
| Qwen3.5-0.8B), greedy, with the server's JSON-schema -> GBNF grammar. |
| |
| ``max_new_tokens`` and :meth:`prompt` are hooks for subclasses (the extraction |
| specialists in :mod:`deg.bench.systems.extractors`). |
| |
| Context: a prompt longer than the model's native context is ``context_exceeded: |
| native``; one that fits the model but not the server's ``n_ctx`` (sized to the GPU |
| memory left beside Clef) is ``context_exceeded: server`` (a harness limit, reported |
| apart). Both score 0.""" |
|
|
| max_new_tokens = MAX_NEW_TOKENS |
|
|
| def __init__(self, url: str | None, label: str | None = None, |
| model_path: str | None = None, quant: str | None = None) -> None: |
| from transformers import AutoTokenizer |
|
|
| from deg.model import base_snapshot |
|
|
| self.url = (url or "http://127.0.0.1:8792").rstrip("/") |
| path = Path(model_path) if model_path else base_snapshot() |
| self.tok = AutoTokenizer.from_pretrained(str(path)) |
| self.native_ctx = native_context(path) |
| self.server_ctx = self._server_ctx() |
| self.name = label or "qwen3.5-0.8b-llamacpp-gbnf" |
| self.info = {"model": (str(path) if model_path else |
| "Qwen3.5-0.8B bf16 GGUF (same base weights, not deg)"), |
| "quantization": quant, "native_context": self.native_ctx, |
| "server_n_ctx": self.server_ctx, |
| "decoding": "greedy, llama.cpp json_schema -> GBNF grammar", |
| "runtime": "llama.cpp server (CUDA, GPU0), batch 1"} |
|
|
| def _server_ctx(self) -> int | None: |
| import urllib.request |
| try: |
| with urllib.request.urlopen(self.url + "/props", timeout=30) as r: |
| props = json.loads(r.read()) |
| return int(props["default_generation_settings"]["n_ctx"]) |
| except Exception: |
| return None |
|
|
| def prompt(self, rec: dict[str, Any]) -> str: |
| return chat_prompt(self.tok, rec) |
|
|
| def predict(self, rec: dict[str, Any]) -> Prediction: |
| import urllib.error |
| import urllib.request |
|
|
| prompt = self.prompt(rec) |
| bos = self.tok.bos_token |
| if bos and prompt.startswith(bos): |
| prompt = prompt[len(bos):] |
| n = len(self.tok(prompt, add_special_tokens=False)["input_ids"]) |
| if self.native_ctx and n >= self.native_ctx: |
| return Prediction(rec["uid"], error=f"context_exceeded: native {n} prompt tokens " |
| f">= {self.native_ctx}", prompt_tokens=n) |
| if self.server_ctx and n >= self.server_ctx: |
| return Prediction(rec["uid"], error=f"context_exceeded: server n_ctx {n} prompt " |
| f"tokens >= {self.server_ctx} (GPU memory)", |
| prompt_tokens=n) |
| last = "" |
| for mode, schema in (("full", grammar_schema(rec["schema"])), |
| ("no-format", grammar_schema(rec["schema"], |
| drop=("x-generate", "format", |
| "pattern")))): |
| body = {"prompt": prompt, "n_predict": self.max_new_tokens, "temperature": 0.0, |
| "top_k": 1, "json_schema": schema, "cache_prompt": False} |
| req = urllib.request.Request(self.url + "/completion", |
| data=json.dumps(body).encode(), |
| headers={"Content-Type": "application/json"}) |
| t0 = time.perf_counter() |
| try: |
| with urllib.request.urlopen(req, timeout=3600) as r: |
| resp = json.loads(r.read()) |
| except urllib.error.HTTPError as exc: |
| last = f"HTTP {exc.code}: {exc.read().decode(errors='replace')[:200]}" |
| if "context" in last.lower(): |
| return Prediction(rec["uid"], error=f"context_exceeded: server {last}"[:240], |
| prompt_tokens=n) |
| continue |
| except (OSError, ValueError) as exc: |
| return Prediction(rec["uid"], error=f"server: {exc}"[:200]) |
| dt = time.perf_counter() - t0 |
| text = resp.get("content", "") |
| gen, ptok = resp.get("tokens_predicted"), resp.get("tokens_evaluated") |
| extra = {"grammar": mode, "server_timings": resp.get("timings")} |
| try: |
| value = json.loads(text) |
| except json.JSONDecodeError as exc: |
| return Prediction(rec["uid"], error=f"json: {exc}"[:200], generated_tokens=gen, |
| prompt_tokens=ptok, latency_s=dt, |
| extra={**extra, "text": text[-300:]}) |
| return Prediction(rec["uid"], value=value, generated_tokens=gen, prompt_tokens=ptok, |
| latency_s=dt, extra=extra) |
| return Prediction(rec["uid"], error=f"grammar: {last}"[:240]) |
|
|
