app.py

import spaces  # must be imported before torch

import html
import json
import sys
import time
from pathlib import Path

import gradio as gr
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

MODEL_ID = "Cloudflare/clef-flash"
MODEL_PATH = Path(snapshot_download(MODEL_ID))
sys.path.insert(0, str(MODEL_PATH))
from joint_schema_model import (  # noqa: E402
    QUESTION_TYPES,
    ClefModel,
    JointSchemaHead,
    collate_records,
    encode_record,
    systemone_answer,
)

# Same as the authors' load_release_model, but with an eager .to("cuda") (ZeroGPU-friendly)
backbone = Qwen3_5ForConditionalGeneration.from_pretrained(
    MODEL_PATH, dtype=torch.bfloat16, attn_implementation="sdpa"
).to("cuda")
backbone.config.use_cache = False
head = JointSchemaHead(**json.loads((MODEL_PATH / "joint_head_config.json").read_text()))
head.load_state_dict(load_file(MODEL_PATH / "joint_head.safetensors"), strict=True)
head = head.to(device="cuda", dtype=torch.bfloat16)
model = ClefModel(backbone, head).eval()
processor = AutoProcessor.from_pretrained(MODEL_PATH)

MAX_IMAGE_SIDE = 1024

# ---------------------------------------------------------------- examples
EX_SUPPORT_STATE = "Our checkout started returning errors and orders are blocked."
EX_SUPPORT_Q = {
    "department": {
        "type": "choice",
        "instructions": "Which team should handle the message?",
        "criteria": {"billing": "Payments or invoices", "technical": "Bugs or outages"},
    },
    "urgency": {"type": "score", "criteria": ["Can wait", "This week", "Today"]},
    "outage": {"type": "noul", "instructions": "Is a service down?"},
}
EX_INVOICE_STATE = {"invoice": {"vendor": "Acme", "total": 1250.0, "currency": "USD", "status": "overdue"}}
EX_INVOICE_Q = {
    "status": {
        "type": "choice",
        "instructions": "What is the invoice status?",
        "criteria": {"paid": "Invoice is paid.", "overdue": "Invoice is past due.", "draft": "Not sent."},
    },
    "large": {"type": "noul", "instructions": "Is the total above 1000 USD?"},
}
EX_REVIEW_STATE = (
    "Ordered the noise-cancelling headphones two weeks ago. Sound is great but the left ear cup "
    "started crackling yesterday and the support chat keeps disconnecting. I want a replacement, not a refund."
)
EX_REVIEW_Q = {
    "sentiment": {
        "type": "score",
        "instructions": "Overall customer sentiment",
        "criteria": ["Very negative", "Negative", "Mixed", "Positive", "Very positive"],
    },
    "action": {
        "type": "choice",
        "instructions": "What should the agent do next?",
        "criteria": {
            "replace": "Ship a replacement unit",
            "refund": "Issue a refund",
            "troubleshoot": "Send troubleshooting steps",
            "escalate": "Escalate to a supervisor",
        },
    },
    "defect": {"type": "noul", "instructions": "Does the customer report a hardware defect?"},
}
EX_FOOD_STATE = {"task": "Moderate a photo uploaded to a restaurant listing.", "listing": "Vegan Garden Bistro"}
EX_FOOD_Q = {
    "category": {
        "type": "choice",
        "instructions": "What does the photo mainly show?",
        "criteria": {"food": "A dish or meal", "interior": "Restaurant interior", "menu": "A menu or text", "people": "People"},
    },
    "matches_listing": {"type": "noul", "instructions": "Is the dish plausibly vegan, consistent with the listing?"},
    "appetizing": {"type": "score", "instructions": "How appetizing is the photo?", "criteria": ["Not at all", "Somewhat", "Very"]},
}
EX_ROOM_STATE = "A guest submitted this photo with a listing for a short-term rental apartment."
EX_ROOM_Q = {
    "room": {
        "type": "choice",
        "instructions": "Which room is shown?",
        "criteria": {"bedroom": "Bedroom", "kitchen": "Kitchen", "living_room": "Living room", "bathroom": "Bathroom"},
    },
    "has_tv": {"type": "noul", "instructions": "Is there a television in the room?"},
    "tidiness": {"type": "score", "instructions": "How tidy is the room?", "criteria": ["Messy", "Average", "Spotless"]},
}


def _j(obj) -> str:
    return json.dumps(obj, indent=2, ensure_ascii=False)


# ---------------------------------------------------------------- rendering
def render_html(questions: dict, answers: dict, probs: dict, latency_ms: float, n_tokens: int) -> str:
    parts = ['<div class="clef-out">']
    for qid, q in questions.items():
        qtype = q["type"]
        p = probs[qid]
        ans = answers[qid]
        if qtype == "noul":
            headline = f"{'true' if p['true'] >= 0.5 else 'false'} &middot; P(true) = {p['true']:.3f}"
            labels = {"true": "true", "false": "false"}
        elif qtype == "choice":
            headline = f"{html.escape(ans['choice'])} &middot; {ans['confidence']:.3f}"
            labels = {k: k for k in p}
        else:
            headline = f"expected score {ans['score']:.2f} / {len(q['criteria']) - 1}"
            labels = {str(i): f"{i} &middot; {html.escape(str(c))}" for i, c in enumerate(q["criteria"])}
        best = max(p, key=p.get)
        instr = html.escape(str(q.get("instructions") or qid))
        parts.append(
            f'<div class="clef-q"><div class="clef-qh"><span class="clef-id">{html.escape(qid)}</span>'
            f'<span class="clef-type">{qtype}</span></div>'
            f'<div class="clef-instr">{instr}</div><div class="clef-head">{headline}</div>'
        )
        order = list(p.keys()) if qtype in ("score", "noul") else sorted(p, key=p.get, reverse=True)
        for opt in order:
            pct = 100 * p[opt]
            cls = "clef-bar best" if opt == best else "clef-bar"
            parts.append(
                f'<div class="clef-row"><div class="clef-lbl">{labels.get(opt, html.escape(opt))}</div>'
                f'<div class="clef-track"><div class="{cls}" style="width:{pct:.1f}%"></div></div>'
                f'<div class="clef-pct">{pct:.1f}%</div></div>'
            )
        parts.append("</div>")
    parts.append(f'<div class="clef-meta">{n_tokens} input tokens &middot; forward pass {latency_ms:.0f} ms</div></div>')
    return "".join(parts)


def _parse_state(state: str):
    s = (state or "").strip()
    if s[:1] in "{[":
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            pass
    return s


def _validate(questions):
    if not isinstance(questions, dict) or not questions:
        raise gr.Error("Questions must be a non-empty JSON object mapping question IDs to questions.")
    for qid, q in questions.items():
        if not isinstance(q, dict) or q.get("type") not in QUESTION_TYPES:
            raise gr.Error(f"{qid}: 'type' must be one of noul, choice, score.")
        if q["type"] == "choice" and not (isinstance(q.get("criteria"), dict) and q["criteria"]):
            raise gr.Error(f"{qid}: a choice question needs 'criteria' as an object of option_id -> description.")
        if q["type"] == "score" and not (isinstance(q.get("criteria"), list) and q["criteria"]):
            raise gr.Error(f"{qid}: a score question needs 'criteria' as a list of level descriptions.")


@spaces.GPU(duration=20)
def decide(state: str, questions_json: str, image=None):
    """Score every allowed option of every typed question about a state with Clef-Flash.

    Args:
        state: the situation to decide on, as free text or a JSON value.
        questions_json: JSON object mapping question IDs to questions. Each question has a `type`
            (`noul` = true/false, `choice` = named options, `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).
        image: optional image that is part of the state.

    Returns:
        An HTML visualisation of per-option probabilities and a SystemOne-style JSON response.
    """
    try:
        questions = json.loads(questions_json)
    except json.JSONDecodeError as e:
        raise gr.Error(f"Questions JSON is invalid: {e}")
    _validate(questions)
    record = {"model": "clef-flash", "state": _parse_state(state), "questions": questions}
    if image is not None:
        image = image.convert("RGB")
        image.thumbnail((MAX_IMAGE_SIDE, MAX_IMAGE_SIDE))
        record["images"] = [image]
    try:
        encoded = encode_record(processor.tokenizer, record, max_length=16384, processor=processor)
    except Exception as e:
        raise gr.Error(f"Could not encode request: {e}")
    batch = collate_records([encoded], processor.tokenizer.pad_token_id, torch.device("cuda"))
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.inference_mode():
        logits = model(batch)[0]
    torch.cuda.synchronize()
    latency_ms = (time.perf_counter() - t0) * 1000
    probs, answers = {}, {}
    for eq, ql in zip(encoded.questions, logits):
        p = dict(zip(eq.option_ids, ql.float().softmax(-1).tolist()))
        probs[eq.question_id] = p
        answers[eq.question_id] = systemone_answer(questions[eq.question_id], p)
    response = {
        "model": "clef-flash",
        "answers": answers,
        "usage": {"input_tokens": len(encoded.input_ids), "output_tokens": 0},
    }
    return render_html(questions, answers, probs, latency_ms, len(encoded.input_ids)), response


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.clef-out { display: flex; flex-direction: column; gap: 12px; }
.clef-q { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 12px 14px; background: var(--block-background-fill); }
.clef-qh { display: flex; gap: 8px; align-items: center; }
.clef-id { font-weight: 700; font-family: var(--font-mono); color: var(--body-text-color); }
.clef-type { font-size: 0.75em; padding: 1px 8px; border-radius: 99px; background: var(--color-accent-soft); color: var(--body-text-color); }
.clef-instr { font-size: 0.9em; opacity: 0.8; margin: 4px 0 6px; color: var(--body-text-color); }
.clef-head { font-weight: 600; margin-bottom: 8px; color: var(--body-text-color); }
.clef-row { display: grid; grid-template-columns: minmax(90px, 38%) 1fr 56px; gap: 8px; align-items: center; margin: 3px 0; }
.clef-lbl { font-size: 0.85em; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; color: var(--body-text-color); }
.clef-track { height: 12px; border-radius: 6px; background: var(--background-fill-secondary); overflow: hidden; }
.clef-bar { height: 100%; background: var(--neutral-400); border-radius: 6px; }
.clef-bar.best { background: var(--color-accent); }
.clef-pct { font-size: 0.85em; text-align: right; font-family: var(--font-mono); color: var(--body-text-color); }
.clef-meta { font-size: 0.8em; opacity: 0.7; color: var(--body-text-color); }
"""

with gr.Blocks(title="Clef-Flash") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Clef-Flash: typed decisions in one forward pass\n"
            "Give [Cloudflare/clef-flash](https://huggingface.co/Cloudflare/clef-flash) a **state** (text, JSON, optionally an image) "
            "and a **schema of typed questions**. It returns a probability for every allowed option of every question "
            "— no text generation, no output parsing. Question types: `noul` (true/false), `choice` (named options), "
            "`score` (ordered levels). [Blog post](https://blog.cloudflare.com/clef-decision-models)"
        )
        with gr.Row():
            with gr.Column(scale=1):
                state = gr.Textbox(
                    label="State (free text or JSON)",
                    lines=5,
                    value=EX_SUPPORT_STATE,
                    placeholder="Describe the situation to decide on, or paste a JSON object…",
                )
                questions = gr.Code(label="Questions schema (JSON)", language="json", value=_j(EX_SUPPORT_Q), lines=14)
                image = gr.Image(label="Image in the state (optional)", type="pil", height=240)
                run = gr.Button("Decide", variant="primary")
            with gr.Column(scale=1):
                out_html = gr.HTML(label="Decisions")
                with gr.Accordion("SystemOne JSON response", open=False):
                    out_json = gr.JSON(label="Response")

        gr.Markdown("### Examples — text / JSON state")
        text_examples = gr.Examples(
            examples=[
                [EX_SUPPORT_STATE, _j(EX_SUPPORT_Q)],
                [_j(EX_INVOICE_STATE), _j(EX_INVOICE_Q)],
                [EX_REVIEW_STATE, _j(EX_REVIEW_Q)],
            ],
            inputs=[state, questions],
            outputs=[out_html, out_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )
        gr.Markdown("### Examples — with an image")
        gr.Examples(
            examples=[
                [_j(EX_FOOD_STATE), _j(EX_FOOD_Q), "examples/burger.jpg"],
                [EX_ROOM_STATE, _j(EX_ROOM_Q), "examples/living_room.jpg"],
            ],
            inputs=[state, questions, image],
            outputs=[out_html, out_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )
        # Text-only examples should not keep a previously loaded image around.
        text_examples.dataset.click(lambda: None, outputs=image, api_visibility="private")

    run.click(decide, inputs=[state, questions, image], outputs=[out_html, out_json], api_name="decide")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)