Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # must be imported before torch | |
| import html | |
| import json | |
| import math | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import gradio as gr | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| MODEL_ID = "nace-ai/drex-v1.5" | |
| # The checkpoint ships its runtime with the weights: code/kev (tokenizer-to-answer stack, Apache-2.0), | |
| # a Qwen3.5 hybrid backbone (bf16 safetensors) and head.pt (pointer head). | |
| CKPT = snapshot_download( | |
| MODEL_ID, | |
| allow_patterns=[ | |
| "*.json", "*.jinja", "inference.py", "head.pt", | |
| "*.safetensors", | |
| "code/*", "code/**/*", | |
| ], | |
| ) | |
| sys.path.insert(0, str(Path(CKPT) / "code")) | |
| from kev.api import SystemOneRequest, output_tokens, to_answers, to_record # noqa: E402 | |
| from kev.checkpoint import load # noqa: E402 | |
| from kev.model import ContextOverflow, SERVE_MAX_BRANCH, SERVE_MAX_STATE # noqa: E402 | |
| # checkpoint.load() defaults to CUDA when available and bf16 there (fp32 otherwise); | |
| # .to("cuda") inside DecisionModel is intercepted by the ZeroGPU hijack and packed at startup. | |
| tok, model = load(CKPT) | |
| print(f"Loaded {MODEL_ID} on {model.device} ({model.dtype}); hybrid={model.hybrid}", flush=True) | |
| MAX_STATE_CHARS = 40000 | |
| EX_REQUEST = { | |
| "state": {"ticket": "I was charged twice for the same order. Please refund the extra payment."}, | |
| "questions": { | |
| "team": { | |
| "type": "choice", | |
| "instructions": "Which team should handle this ticket?", | |
| "criteria": {"billing": "Payments, charges, and refunds", "technical": "Bugs and outages", "other": "Anything else"}, | |
| }, | |
| "refund": {"type": "noul", "instructions": "Does the customer explicitly ask for a refund?"}, | |
| "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["Routine", "Soon", "Urgent"]}, | |
| }, | |
| } | |
| EX_SERVICE = { | |
| "state": {"alert": "checkout-api error_rate jumped from 0.2% to 34% five minutes ago; latency_p99 is 4.8 s and orders are failing."}, | |
| "questions": { | |
| "sev": { | |
| "type": "score", | |
| "instructions": "How severe is this incident?", | |
| "criteria": ["Minor", "Elevated", "Major", "Critical"], | |
| }, | |
| "page": {"type": "noul", "instructions": "Should the on-call engineer be paged right now?"}, | |
| "area": { | |
| "type": "choice", | |
| "instructions": "Which area is most likely involved?", | |
| "criteria": { | |
| "payments": "Charging, refunds, invoices", | |
| "frontend": "Checkout UI and browser errors", | |
| "backend": "API failures and latency", | |
| "infra": "Deployments, scaling, cloud resources", | |
| }, | |
| }, | |
| }, | |
| } | |
| EX_REVIEW = { | |
| "state": {"review": "The room was spotless and the staff were wonderful, but the room heater rattled all night and I barely slept.", "stars": 2}, | |
| "questions": { | |
| "sentiment": { | |
| "type": "score", | |
| "instructions": "How positive is this review overall?", | |
| "criteria": ["Very negative", "Negative", "Neutral", "Positive", "Very positive"], | |
| }, | |
| "would_return": {"type": "noul", "instructions": "Does the reviewer sound likely to return?"}, | |
| "followup": { | |
| "type": "choice", | |
| "instructions": "What should the hotel do first?", | |
| "criteria": { | |
| "maintenance": "Fix the heater", | |
| "compensation": "Offer a refund or discount", | |
| "thanks": "Thank the reviewer", | |
| "nothing": "No action needed", | |
| }, | |
| }, | |
| }, | |
| } | |
| def _dump(value) -> str: | |
| return value if isinstance(value, str) else json.dumps(value, indent=2, ensure_ascii=False) | |
| def _first_error(error) -> str: | |
| """Pydantic wraps every failing question in one long message; keep the head of it.""" | |
| text = str(error).replace("\n", " ") | |
| return text[:300] + ("…" if len(text) > 300 else "") | |
| def _parse_state(state: str): | |
| """Plain text passes through; JSON-looking text is parsed into a structured state.""" | |
| text = (state or "").strip() | |
| if not text: | |
| raise gr.Error("Please provide a state (plain text or JSON).") | |
| if text[0] in "{[": | |
| try: | |
| return json.loads(text) | |
| except json.JSONDecodeError: | |
| pass | |
| return text | |
| def _gpu_duration(state: str, questions_json: str, *args, **kwargs) -> int: | |
| # Measured on ZeroGPU: 87-token sample ~3.6 s cold / ~0.1 s warm; 8k-token state with | |
| # 3 questions ~4.7 s model time. Scale with input size, capped well under the tier limit. | |
| chars = len(state or "") + len(questions_json or "") | |
| return int(min(60, 10 + chars / 3000)) | |
| def decide(state: str, questions_json: str): | |
| """Answer typed questions about a document with Drex v1.5. | |
| Args: | |
| state: The document to decide about. Plain text, or JSON (object/list) which is | |
| rendered as indented text for the model. | |
| questions_json: JSON object mapping question names to questions. Each question has | |
| a `type` — `choice` (named options), `noul` (yes/no) or `score` (ordered levels) — | |
| optional `instructions`, and `criteria` (object of option → description for | |
| choice, list of level descriptions for score, optional true/false descriptions | |
| for noul). | |
| Returns: | |
| Per-question probability distributions rendered as bars, plus the SystemOne JSON | |
| response (answer, confidence and probability per option for every question). | |
| """ | |
| parsed_state = _parse_state(state) | |
| try: | |
| parsed_questions = json.loads(questions_json or "") | |
| except json.JSONDecodeError as error: | |
| raise gr.Error(f"Questions must be valid JSON: {error}") | |
| if not isinstance(parsed_questions, dict) or not parsed_questions: | |
| raise gr.Error("Questions must be a non-empty JSON object keyed by question name.") | |
| request = {"state": parsed_state, "questions": parsed_questions} | |
| if len(json.dumps(parsed_state, ensure_ascii=False)) > MAX_STATE_CHARS: | |
| raise gr.Error(f"State is too large (>{MAX_STATE_CHARS} characters). Use a shorter document.") | |
| try: | |
| req = SystemOneRequest.model_validate(request) | |
| except Exception as error: | |
| raise gr.Error(f"Invalid request: {_first_error(error)}") | |
| try: | |
| rec, meta = to_record(req) | |
| enc = model.encode(tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=True) | |
| except ContextOverflow as error: | |
| raise gr.Error(f"Too long for the context window: {error}") | |
| torch.cuda.synchronize() | |
| started = time.perf_counter() | |
| with torch.no_grad(): | |
| probs = [p.float().tolist() for p in model.probs(enc)] | |
| torch.cuda.synchronize() | |
| latency_ms = (time.perf_counter() - started) * 1000 | |
| if len(probs) != len(meta) or any( | |
| len(row) != len(question["keys"]) or any(not math.isfinite(value) for value in row) | |
| for row, question in zip(probs, meta) | |
| ): | |
| raise gr.Error("Model readout is incomplete or non-finite; try a shorter request.") | |
| answers = to_answers(probs, meta) | |
| response = { | |
| "model": "drex-v1.5", | |
| "answers": answers, | |
| "usage": { | |
| "input_tokens": len(enc["ids"]), | |
| "output_tokens": output_tokens(tok, answers), | |
| }, | |
| "latency_ms": round(latency_ms, 1), | |
| } | |
| return _render(parsed_questions, answers, latency_ms, len(enc["ids"])), response | |
| def _render(questions: dict, answers: dict, latency_ms: float, n_tokens: int) -> str: | |
| cards = [] | |
| for qid, question in questions.items(): | |
| ans = answers[qid] | |
| if ans["type"] == "noul": | |
| options = [("yes", "yes", ans["noul"]), ("no", "no", 1.0 - ans["noul"])] | |
| elif ans["type"] == "choice": | |
| criteria = question.get("criteria") or {} | |
| options = [ | |
| (k, f"{k} — {v}" if isinstance(v, str) and v else k, ans["probabilities"][k]) | |
| for k, v in criteria.items() | |
| ] | |
| else: # score | |
| legend = ans["legend"] | |
| options = [(f"{i}: {legend[i]}", legend[i], ans["probabilities"][str(i)]) for i in range(len(legend))] | |
| best = max(options, key=lambda o: o[2])[0] | |
| rows = [] | |
| for key, label, value in options: | |
| highlight = "drex-best" if key == best else "" | |
| rows.append( | |
| f'<div class="drex-row {highlight}"><div class="drex-label">{html.escape(str(label))}</div>' | |
| f'<div class="drex-bar"><div class="drex-fill" style="width:{value * 100:.1f}%"></div></div>' | |
| f'<div class="drex-val">{value * 100:.1f}%</div></div>' | |
| ) | |
| extra = "" | |
| if ans["type"] == "noul": | |
| extra = f' · P(yes) <b>{ans["noul"] * 100:.1f}%</b>' | |
| elif ans["type"] == "score": | |
| extra = f' · expected score <b>{ans["score"]:.2f}</b> of {len(ans["legend"]) - 1}' | |
| if "confidence" in ans: | |
| extra += f' · confidence <b>{ans["confidence"] * 100:.1f}%</b>' | |
| instr = question.get("instructions") or qid | |
| cards.append( | |
| f'<div class="drex-card"><div class="drex-head"><span class="drex-id">{html.escape(str(qid))}</span>' | |
| f'<span class="drex-type">{ans["type"]}</span></div>' | |
| f'<div class="drex-instr">{html.escape(str(instr))}</div>' | |
| f'<div class="drex-answer">→ <b>{html.escape(str(best))}</b>{extra}</div>{"".join(rows)}</div>' | |
| ) | |
| footer = f'<div class="drex-foot">{n_tokens} input tokens · model time {latency_ms:.0f} ms · no text generated</div>' | |
| return '<div class="drex-wrap">' + "".join(cards) + footer + "</div>" | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .drex-wrap { display: flex; flex-direction: column; gap: 12px; } | |
| .drex-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 12px 14px; } | |
| .drex-head { display: flex; gap: 8px; align-items: center; margin-bottom: 2px; } | |
| .drex-id { font-weight: 700; font-family: var(--font-mono); } | |
| .drex-type { font-size: 0.75em; padding: 1px 8px; border-radius: 999px; background: var(--color-accent-soft); } | |
| .drex-instr { opacity: 0.8; font-size: 0.92em; } | |
| .drex-answer { margin: 6px 0 8px; } | |
| .drex-row { display: grid; grid-template-columns: minmax(120px, 40%) 1fr 56px; gap: 8px; align-items: center; margin: 3px 0; font-size: 0.9em; } | |
| .drex-label { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } | |
| .drex-bar { height: 12px; border-radius: 6px; background: var(--block-background-fill); border: 1px solid var(--border-color-primary); overflow: hidden; } | |
| .drex-fill { height: 100%; background: var(--neutral-400); } | |
| .drex-best .drex-fill { background: var(--color-accent); } | |
| .drex-best .drex-label { font-weight: 600; } | |
| .drex-val { text-align: right; font-family: var(--font-mono); } | |
| .drex-foot { font-size: 0.8em; opacity: 0.7; } | |
| """ | |
| with gr.Blocks(title="Drex v1.5") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# Drex v1.5 — typed decisions with probabilities\n" | |
| "[nace-ai/drex-v1.5](https://huggingface.co/nace-ai/drex-v1.5) is a 9B **decision model** from " | |
| "[Nace.AI](https://www.nace.ai/): give it a **state** (a document: text or JSON) and named typed questions — " | |
| "`choice` (pick an option), `noul` (yes/no) or `score` (ordered levels) — and it returns a " | |
| "probability for every option, one forward pass per question. No text is generated. " | |
| "Built on [MiMo-V2.6-Distill-Qwen-9B](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Distill-Qwen-9B) " | |
| "(Qwen3.5 hybrid backbone) with a pointer-head readout." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| state = gr.Textbox( | |
| label="State (the document: plain text or JSON)", | |
| lines=5, | |
| placeholder="Paste a support ticket, a log line, a JSON record, a contract clause…", | |
| ) | |
| questions = gr.Code( | |
| label="Questions (JSON: name → {type, instructions, criteria})", | |
| language="json", | |
| value=_dump(EX_REQUEST["questions"]), | |
| lines=16, | |
| ) | |
| run = gr.Button("Decide", variant="primary") | |
| with gr.Column(scale=1): | |
| result_html = gr.HTML(label="Decisions") | |
| with gr.Accordion("SystemOne JSON response", open=False): | |
| result_json = gr.JSON(label="Response") | |
| run.click( | |
| decide, | |
| inputs=[state, questions], | |
| outputs=[result_html, result_json], | |
| api_name="decide", | |
| ) | |
| gr.Examples( | |
| label="Examples", | |
| examples=[ | |
| [_dump(EX_REQUEST["state"]), _dump(EX_REQUEST["questions"])], | |
| [_dump(EX_SERVICE["state"]), _dump(EX_SERVICE["questions"])], | |
| [_dump(EX_REVIEW["state"]), _dump(EX_REVIEW["questions"])], | |
| ], | |
| inputs=[state, questions], | |
| outputs=[result_html, result_json], | |
| fn=decide, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) |