Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("D3_FAST", "0") | |
| import spaces # must come before torch / transformers | |
| import html | |
| import json | |
| import time | |
| from pathlib import Path | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoModel | |
| try: | |
| import spaces | |
| def _gpu_placeholder(): | |
| return 0 | |
| except ImportError: | |
| pass | |
| MODEL_ID = "vllm-sr/d3-mini" | |
| HERE = Path(__file__).parent | |
| IMAGE_EXT = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"} | |
| VIDEO_EXT = {".mp4", ".webm", ".mov", ".mkv", ".avi"} | |
| # The d3 runtime places weights itself; load on CPU, then move to CUDA so ZeroGPU can intercept the move. | |
| # flash-linear-attention is deliberately NOT installed: its kernels are GPU-only and would block | |
| # the CPU load; transformers falls back to its (correct) reference implementation instead. | |
| model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, device="cpu") | |
| # model.to("cuda") | |
| # --------------------------------------------------------------------------- examples (from the model card) | |
| TEXT_STATE = ( | |
| "The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today." | |
| ) | |
| TEXT_QUESTIONS = { | |
| "route": { | |
| "type": "choice", | |
| "instructions": "Which team should handle this request?", | |
| "criteria": { | |
| "returns": "Refunds, replacements and damaged deliveries", | |
| "billing": "Payments, invoices and charges", | |
| "technical": "Product setup and faults", | |
| }, | |
| }, | |
| "receipt": {"type": "noul", "instructions": "Does the customer have a receipt?"}, | |
| "urgency": { | |
| "type": "score", | |
| "instructions": "How urgent is this request?", | |
| "criteria": ["Routine", "Soon", "Today"], | |
| }, | |
| } | |
| IMAGE_STATE = "The customer says the blender arrived cracked and attached the receipt." | |
| IMAGE_QUESTIONS = { | |
| "route": TEXT_QUESTIONS["route"], | |
| "on_receipt": {"type": "noul", "instructions": "Does the receipt list the blender?"}, | |
| "payment": { | |
| "type": "choice", | |
| "instructions": "How was the order paid?", | |
| "criteria": {"card": None, "cash": None, "gift card": None}, | |
| }, | |
| } | |
| VIDEO_STATE = "A short clip from a test camera." | |
| VIDEO_QUESTIONS = { | |
| "direction": { | |
| "type": "choice", | |
| "instructions": "Which way does the square move?", | |
| "criteria": { | |
| "right": "From left to right", | |
| "left": "From right to left", | |
| "still": "It does not move", | |
| }, | |
| }, | |
| "color_change": {"type": "noul", "instructions": "Does the square change color?"}, | |
| } | |
| JSON_STATE = json.dumps( | |
| { | |
| "ticket_id": "T-4821", | |
| "channel": "email", | |
| "message": "I was charged twice for my subscription this month. Please fix it ASAP, I'm about to cancel.", | |
| "customer_tier": "premium", | |
| }, | |
| indent=2, | |
| ) | |
| JSON_QUESTIONS = { | |
| "route": TEXT_QUESTIONS["route"], | |
| "churn_risk": { | |
| "type": "noul", | |
| "instructions": "Is the customer at risk of cancelling?", | |
| }, | |
| "sentiment": { | |
| "type": "score", | |
| "instructions": "How satisfied is the customer?", | |
| "criteria": ["Very unhappy", "Unhappy", "Neutral", "Happy", "Very happy"], | |
| }, | |
| } | |
| def dumps(obj) -> str: | |
| return json.dumps(obj, indent=2, ensure_ascii=False) | |
| TEMPLATES = { | |
| "choice": { | |
| "type": "choice", | |
| "instructions": "Which option fits best?", | |
| "criteria": {"option_a": "Description of A", "option_b": "Description of B"}, | |
| }, | |
| "noul": {"type": "noul", "instructions": "Is this statement true?"}, | |
| "score": { | |
| "type": "score", | |
| "instructions": "Rate this on a scale.", | |
| "criteria": ["Low", "Medium", "High"], | |
| }, | |
| } | |
| def add_question(current: str, kind: str) -> str: | |
| """Append a template question of the given kind to the questions JSON.""" | |
| try: | |
| questions = json.loads(current) if current.strip() else {} | |
| if not isinstance(questions, dict): | |
| raise ValueError | |
| except ValueError: | |
| raise gr.Error("Fix the questions JSON before adding a new question.") | |
| base = {"choice": "my_choice", "noul": "my_yes_no", "score": "my_score"}[kind] | |
| name, i = base, 2 | |
| while name in questions: | |
| name, i = f"{base}_{i}", i + 1 | |
| questions[name] = json.loads(json.dumps(TEMPLATES[kind])) | |
| return dumps(questions) | |
| # --------------------------------------------------------------------------- rendering | |
| def _bar(label: str, p: float, best: bool, sub: str = "") -> str: | |
| pct = max(0.0, min(1.0, p)) * 100 | |
| color = "#f97316" if best else "#a78bfa" | |
| sub_html = f'<span class="d3-sub">{html.escape(sub)}</span>' if sub else "" | |
| return ( | |
| '<div class="d3-row">' | |
| f'<div class="d3-label">{"<b>" if best else ""}{html.escape(label)}{"</b>" if best else ""}{sub_html}</div>' | |
| f'<div class="d3-track"><div class="d3-fill" style="width:{pct:.1f}%;background:{color}"></div></div>' | |
| f'<div class="d3-pct">{pct:.1f}%</div>' | |
| "</div>" | |
| ) | |
| def render_answers(questions: dict, answers: dict, elapsed: float, tokens: int) -> str: | |
| cards = [] | |
| for key, ans in answers.items(): | |
| q = questions.get(key, {}) | |
| title = html.escape(str(q.get("instructions") or key)) | |
| kind = ans.get("type", q.get("type", "")) | |
| badge = {"choice": "Choice", "noul": "Yes / No", "score": "Score"}.get(kind, kind) | |
| head = f'<div class="d3-head"><span class="d3-badge">{badge}</span><code>{html.escape(key)}</code> {title}</div>' | |
| if "error" in ans: | |
| body = f'<div class="d3-err">{html.escape(ans["error"])}: {html.escape(ans.get("message", ""))}</div>' | |
| elif kind == "noul": | |
| p = float(ans["noul"]) | |
| verdict = "Yes" if p >= 0.5 else "No" | |
| body = ( | |
| f'<div class="d3-verdict">{verdict}</div>' | |
| + _bar("Yes", p, p >= 0.5) | |
| + _bar("No", 1 - p, p < 0.5) | |
| ) | |
| elif kind == "choice": | |
| probs = ans["probabilities"] | |
| crit = q.get("criteria") or {} | |
| best = ans["choice"] | |
| body = ( | |
| f'<div class="d3-verdict">{html.escape(best)}' | |
| f'<span class="d3-conf">confidence {ans.get("confidence", 0):.2f}</span></div>' | |
| ) | |
| for k, p in sorted(probs.items(), key=lambda kv: -kv[1]): | |
| desc = crit.get(k) if isinstance(crit, dict) else None | |
| body += _bar(k, p, k == best, desc if isinstance(desc, str) else "") | |
| elif kind == "score": | |
| probs = ans["probabilities"] | |
| legend = ans.get("legend", {}) | |
| n = len(probs) | |
| best = max(probs, key=probs.get) | |
| body = ( | |
| f'<div class="d3-verdict">{ans["score"]:.2f} <span class="d3-conf">expected level on 0–{n - 1}' | |
| f' · confidence {ans.get("confidence", 0):.2f}</span></div>' | |
| ) | |
| for k, p in probs.items(): | |
| body += _bar(f"{k}: {legend.get(k, k)}", p, k == best) | |
| else: | |
| body = f"<pre>{html.escape(dumps(ans))}</pre>" | |
| cards.append(f'<div class="d3-card">{head}{body}</div>') | |
| footer = f'<div class="d3-foot">{len(answers)} question(s) · {tokens} input tokens · {elapsed * 1000:.0f} ms on GPU</div>' | |
| return '<div class="d3-wrap">' + "".join(cards) + footer + "</div>" | |
| # --------------------------------------------------------------------------- inference | |
| def _paths(files) -> list[str]: | |
| out = [] | |
| for f in files or []: | |
| p = f if isinstance(f, str) else getattr(f, "name", None) or (f.get("path") if isinstance(f, dict) else None) | |
| if p: | |
| out.append(str(p)) | |
| return out | |
| def _gpu_duration(state, questions, images=None, videos=None): | |
| """Seconds to reserve, sized from measurements on this Space's ZeroGPU worker. | |
| Measured (48 GB `large`, torch-reference kernels, no fla): | |
| * warm text request (3 questions, 345 tokens): 68 ms of GPU work | |
| * warm 1-image request (3 questions, 2938 tokens): 551 ms | |
| * warm 1-video request (2 questions, 1720 tokens): 288 ms | |
| * the first request after boot: ~5.5 s wall — one-time first-touch cost of | |
| streaming the packed 9.08 GB of bf16 weights into VRAM. | |
| The floor of 8 s covers that first-touch (the scheduler recycles idle workers, | |
| so a fresh one shows up occasionally) plus a small margin; the per-input terms | |
| scale for big states, image stacks and long videos (decoding happens inside | |
| the GPU reservation too). A request that overruns its reservation only logs a | |
| warning, so the estimate stays lean: a smaller number ranks higher in the queue | |
| and costs every visitor less quota. | |
| """ | |
| chars = len(json.dumps(state, ensure_ascii=False, default=str)) + len( | |
| json.dumps(questions, ensure_ascii=False, default=str) | |
| ) | |
| seconds = 8 + chars / 300 + 2 * len(images or []) + (5 if videos else 0) | |
| return int(min(30, seconds)) | |
| # @spaces.GPU(duration=_gpu_duration) | |
| def _run(state, questions, images, videos): | |
| t0 = time.perf_counter() | |
| with torch.inference_mode(): | |
| result = model.system_one(state=state, questions=questions, images=images or None, videos=videos or None) | |
| return result, time.perf_counter() - t0 | |
| def decide(state: str, questions_json: str, media: list | None = None): | |
| """Answer Choice / Yes-No / Score questions about a text or JSON state, optionally with images and videos. | |
| Args: | |
| state: the input to decide about, plain text or a JSON document. | |
| questions_json: JSON object mapping question IDs to {"type": "choice"|"noul"|"score", | |
| "instructions": str, "criteria": ...}. Choice criteria map option keys to descriptions | |
| (or null); score criteria are an ordered list of 2-10 levels; noul is a yes/no question. | |
| media: optional list of image (png/jpg/webp) and video (mp4/webm/mov/mkv) file paths that every | |
| question sees. | |
| Returns: | |
| An HTML rendering of the answers and the raw System One response (answers with probabilities). | |
| """ | |
| if not state or not state.strip(): | |
| raise gr.Error("Enter a state (text or JSON) to decide about.") | |
| try: | |
| questions = json.loads(questions_json) | |
| except json.JSONDecodeError as e: | |
| raise gr.Error(f"Questions are not valid JSON: {e}") | |
| if not isinstance(questions, dict) or not questions: | |
| raise gr.Error("Questions must be a non-empty JSON object of {id: question}.") | |
| stripped = state.strip() | |
| parsed_state = state | |
| if stripped[:1] in "{[": | |
| try: | |
| parsed_state = json.loads(stripped) | |
| except json.JSONDecodeError: | |
| parsed_state = state | |
| images, videos = [], [] | |
| for p in _paths(media): | |
| ext = Path(p).suffix.lower() | |
| if ext in IMAGE_EXT: | |
| images.append(p) | |
| elif ext in VIDEO_EXT: | |
| videos.append(p) | |
| else: | |
| raise gr.Error(f"Unsupported file type: {Path(p).name}") | |
| try: | |
| result, elapsed = _run(parsed_state, questions, images, videos) | |
| except ValueError as e: | |
| raise gr.Error(str(e)) | |
| tokens = result.get("usage", {}).get("input_tokens", 0) | |
| return render_answers(questions, result["answers"], elapsed, tokens), result | |
| # --------------------------------------------------------------------------- UI | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .d3-wrap { display: flex; flex-direction: column; gap: 12px; } | |
| .d3-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 12px 14px; background: var(--block-background-fill); } | |
| .d3-head { font-size: 0.95em; margin-bottom: 6px; color: var(--body-text-color); } | |
| .d3-head code { margin: 0 6px; } | |
| .d3-badge { font-size: 0.75em; padding: 2px 8px; border-radius: 999px; background: #fde68a; color: #78350f; font-weight: 600; } | |
| .d3-verdict { font-size: 1.35em; font-weight: 700; margin: 4px 0 8px; color: var(--body-text-color); } | |
| .d3-conf { font-size: 0.6em; font-weight: 400; margin-left: 10px; opacity: 0.7; } | |
| .d3-row { display: grid; grid-template-columns: minmax(120px, 38%) 1fr 56px; gap: 8px; align-items: center; margin: 3px 0; } | |
| .d3-label { font-size: 0.9em; color: var(--body-text-color); overflow-wrap: anywhere; } | |
| .d3-sub { display: block; font-size: 0.8em; opacity: 0.65; } | |
| .d3-track { height: 12px; border-radius: 6px; background: var(--background-fill-secondary); overflow: hidden; } | |
| .d3-fill { height: 100%; border-radius: 6px; } | |
| .d3-pct { font-variant-numeric: tabular-nums; font-size: 0.85em; text-align: right; color: var(--body-text-color); } | |
| .d3-err { color: #dc2626; } | |
| .d3-foot { font-size: 0.8em; opacity: 0.7; color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="d3-mini Decision") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# d3-mini: multimodal decision model\n" | |
| "Give [vllm-sr/d3-mini](https://huggingface.co/vllm-sr/d3-mini) (4.54B) an input — text or JSON, " | |
| "plus optional images and videos — and any number of **Choice**, **Yes / No** and **Score** questions. " | |
| "It answers them all in one call with a probability for every option, without generating text. " | |
| "Part of [vLLM Semantic Router](https://github.com/vllm-project/semantic-router)." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| state = gr.Textbox( | |
| label="State (text or JSON)", | |
| lines=5, | |
| placeholder="Describe the situation, paste a ticket, or a JSON document…", | |
| ) | |
| media = gr.File( | |
| label="Images / videos (optional, several allowed)", | |
| file_count="multiple", | |
| file_types=sorted(IMAGE_EXT | VIDEO_EXT), | |
| ) | |
| questions = gr.Code( | |
| label="Questions (JSON)", | |
| language="json", | |
| value=dumps(TEXT_QUESTIONS), | |
| lines=16, | |
| ) | |
| with gr.Row(): | |
| add_choice = gr.Button("+ Choice", size="sm") | |
| add_noul = gr.Button("+ Yes / No", size="sm") | |
| add_score = gr.Button("+ Score", size="sm") | |
| run = gr.Button("Decide", variant="primary") | |
| with gr.Column(scale=1): | |
| answers_html = gr.HTML(label="Answers") | |
| with gr.Accordion("Raw response (JSON)", open=False): | |
| raw = gr.JSON(label="System One response") | |
| with gr.Accordion("Question format", open=False): | |
| gr.Markdown( | |
| "- `choice`: `criteria` maps option keys to descriptions (or `null`), 1–255 options.\n" | |
| "- `noul`: yes / no; optional `criteria` may describe `\"true\"` and `\"false\"`.\n" | |
| "- `score`: `criteria` is an ordered list of 2–10 levels; the answer is the expected level.\n\n" | |
| "Images are read at up to 1.6 MP; videos at 2 fps (max 32 frames, 0.2 MP each)." | |
| ) | |
| add_choice.click(lambda q: add_question(q, "choice"), questions, questions, api_visibility="private") | |
| add_noul.click(lambda q: add_question(q, "noul"), questions, questions, api_visibility="private") | |
| add_score.click(lambda q: add_question(q, "score"), questions, questions, api_visibility="private") | |
| run.click(decide, [state, questions, media], [answers_html, raw], api_name="decide") | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) |