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
@spaces.GPU(duration=1)
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'{html.escape(sub)}' if sub else ""
return (
''
f'
{"" if best else ""}{html.escape(label)}{"" if best else ""}{sub_html}
'
f'
'
f'
{pct:.1f}%
'
"
{badge}{html.escape(key)} {title}
'
if "error" in ans:
body = f'{html.escape(ans["error"])}: {html.escape(ans.get("message", ""))}
'
elif kind == "noul":
p = float(ans["noul"])
verdict = "Yes" if p >= 0.5 else "No"
body = (
f'{verdict}
'
+ _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'{html.escape(best)}'
f'confidence {ans.get("confidence", 0):.2f}
'
)
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'{ans["score"]:.2f} expected level on 0–{n - 1}'
f' · confidence {ans.get("confidence", 0):.2f}
'
)
for k, p in probs.items():
body += _bar(f"{k}: {legend.get(k, k)}", p, k == best)
else:
body = f"{html.escape(dumps(ans))}"
cards.append(f'{head}{body}
')
footer = f''
return '' + "".join(cards) + footer + "
"
# --------------------------------------------------------------------------- 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)