Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Fix blank kind cells in examples (Dropdown rejects empty string); allow custom value
54a9672 verified | """JEV-9B — typed decisions in one forward pass (System 1) + System 2 escalation. | |
| Interactive demo for https://huggingface.co/autotrust/JEV-9B — AutoTrust's integrated | |
| System 1 + System 2 model built with the Blocks-of-Experts recipe on a frozen | |
| Qwen3.5-9B backbone. System 1 answers typed questions (`noul` yes/no, `choice` | |
| over 2-16 options, `score` on a 0-5 scale) in one prefill pass with a calibrated | |
| probability distribution; System 2 is the pristine base model (adapter off) and can | |
| be escalated to for deliberate step-by-step reasoning. | |
| """ | |
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False") | |
| import spaces # noqa: E402 -- MUST precede torch / transformers / any CUDA-touching import | |
| import html # noqa: E402 | |
| import json # noqa: E402 | |
| import time # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import torch # noqa: E402 # noqa: F401 (imported for dtype selection downstream) | |
| import jev_core as jc # noqa: E402 | |
| N_QUESTIONS = 4 | |
| LETTERS = jc.LETTERS | |
| KINDS = list(jc.KINDS) | |
| SCORE_LEVELS = [str(i) for i in range(6)] | |
| TOPIC_OPTIONS = [ | |
| "open-source software", "programming languages & tools", "AI / machine learning", | |
| "security & privacy", "hardware & retro computing", "science & health", | |
| "automotive & energy", "politics & government", "business & finance", "culture & lifestyle", | |
| ] | |
| INTENT_OPTIONS = [ | |
| "asking a question", "sharing personal experience", "providing factual information", | |
| "arguing an opinion", "attacking another commenter", "moderating the discussion", | |
| ] | |
| print( | |
| f"[jev-9b] loaded {jc.REPO}: backbone dtype={next(jc._BACKBONE.parameters()).dtype}, " | |
| f"head dtype={jc._HEAD.weight.dtype}, temperatures={jc.TEMPERATURES}", | |
| flush=True, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def _parse_options(text: str) -> list: | |
| """Options are comma-separated; newlines are accepted as separators too.""" | |
| if not text: | |
| return [] | |
| return [o.strip() for o in text.replace("\n", ",").split(",") if o.strip()] | |
| def _build_questions(fields) -> list: | |
| """Turn the (type, instructions, options) UI rows into jev_core.decide calls.""" | |
| questions = [] | |
| for i in range(N_QUESTIONS): | |
| kind = (fields[3 * i] or "").strip() | |
| instructions = (fields[3 * i + 1] or "").strip() | |
| raw_options = fields[3 * i + 2] or "" | |
| if not instructions: | |
| continue # a row with no instructions is skipped | |
| if kind not in KINDS: | |
| raise gr.Error(f"Question {i + 1}: unknown question type {kind!r}.") | |
| if kind == "choice": | |
| options = _parse_options(raw_options) | |
| if not 2 <= len(options) <= 16: | |
| raise gr.Error( | |
| f"Question {i + 1} is a `choice` question, which needs 2-16 " | |
| f"comma-separated options; got {len(options)}." | |
| ) | |
| elif kind == "score": | |
| options = SCORE_LEVELS # fixed 0-5 levels; the instructions describe the scale | |
| else: # noul | |
| options = ["false", "true"] | |
| questions.append({"kind": kind, "instructions": instructions, "options": options}) | |
| if not questions: | |
| raise gr.Error("Add at least one question — give it some instructions text.") | |
| return questions | |
| def _bars_html(items: list, winner) -> str: | |
| """Render labelled probability bars; the winner is highlighted.""" | |
| out = [] | |
| for label, prob in items: | |
| hot = label == winner | |
| out.append( | |
| '<div style="display:flex;align-items:center;gap:10px;margin:2px 0">' | |
| f'<span style="flex:0 0 230px;text-align:right;font-size:12.5px;' | |
| f'color:var(--body-text-color);opacity:{1.0 if hot else 0.6};' | |
| f'overflow:hidden;text-overflow:ellipsis;white-space:nowrap">' | |
| f"{html.escape(str(label))}</span>" | |
| '<div style="flex:1;background:var(--border-color-primary);height:13px;border-radius:3px">' | |
| f'<div style="width:{max(prob * 100, 0.6):.2f}%;height:13px;border-radius:3px;' | |
| f'background:var(--color-accent);opacity:{1.0 if hot else 0.45}"></div></div>' | |
| f'<span style="flex:0 0 58px;font-size:12.5px;opacity:{1.0 if hot else 0.6}">' | |
| f"{prob * 100:.1f}%</span></div>" | |
| ) | |
| return "".join(out) | |
| def _answer_html(q: dict, probs: dict) -> str: | |
| """One question -> headline answer plus a labelled probability bar per option.""" | |
| kind, head_txt = q["kind"], html.escape(q["instructions"]) | |
| items = list(probs.items()) | |
| if kind == "choice": | |
| winner = max(probs, key=probs.get) | |
| headline = html.escape(str(winner)) | |
| note = f"confidence {probs[winner]:.3f}" | |
| elif kind == "score": | |
| ev = sum(int(k) * v for k, v in probs.items()) | |
| winner = max(probs, key=probs.get) | |
| headline = f"expected level {ev:.2f} / 5" | |
| note = f"most likely level {winner} ({probs[winner] * 100:.1f}%)" | |
| items = [(f"level {k}", v) for k, v in items] | |
| else: # noul | |
| p_yes = probs["true"] | |
| winner = "true" if p_yes >= 0.5 else "false" | |
| headline = "true" if p_yes >= 0.5 else "false" | |
| note = f"P(true) = {p_yes:.3f}" | |
| return ( | |
| '<div style="padding:10px 0;border-bottom:1px solid var(--border-color-primary)">' | |
| f'<div style="font-size:12.5px;opacity:0.75;margin-bottom:3px">' | |
| f'<code style="font-size:11.5px">{kind}</code> {head_txt}</div>' | |
| f'<div style="font-size:19px;font-weight:600;margin:2px 0">{headline} ' | |
| f'<span style="font-weight:400;font-size:13px;opacity:0.75">{note}</span></div>' | |
| f"{_bars_html(items, winner)}</div>" | |
| ) | |
| def _split_reasoning(text: str) -> tuple: | |
| """Split a System 2 reply into (reasoning, answer) if it carries think tags.""" | |
| open_tag, close_tag = "<think>", "</think>" | |
| if close_tag in text: | |
| reasoning = text.split(close_tag)[0].replace(open_tag, "").strip() | |
| answer = text.split(close_tag, 1)[1].strip() | |
| return reasoning, answer | |
| return "", text.strip() | |
| def _matched_option(answer: str, options: list): | |
| """Return the option the System 2 answer's last line picked, if any.""" | |
| if not answer.strip(): | |
| return None | |
| last = answer.strip().splitlines()[-1].strip().rstrip(".").strip("`").lower() | |
| for i, o in enumerate(options): | |
| ol = str(o).lower() | |
| if last == ol or last == f"{LETTERS[i]}) {ol}" or last == LETTERS[i].lower(): | |
| return o | |
| for o in options: # fall back to a contained mention | |
| if str(o).lower() in last: | |
| return o | |
| return None | |
| # --------------------------------------------------------------------------- | |
| # GPU handlers | |
| # --------------------------------------------------------------------------- | |
| def decide(state: str, *fields): | |
| """Answer several typed questions about one piece of state text in one forward pass each. | |
| Args: | |
| state: the context to reason about — a customer message, a forum post, a diff. | |
| fields: four (type, instructions, options) triples. `type` is choice / score / noul; | |
| `options` is a comma-separated list (used by `choice`, ignored by the others). | |
| Returns: | |
| (answers as HTML, a one-line status, the raw answers as JSON) | |
| """ | |
| state = (state or "").strip() | |
| if not state: | |
| raise gr.Error("Give the model some state text to reason about.") | |
| questions = _build_questions([f if isinstance(f, str) else "" for f in fields]) | |
| state_used, truncated, n_tokens = jc.truncate_state(state) | |
| t0 = time.perf_counter() | |
| answers = [jc.decide(q["kind"], state_used, q["instructions"], q["options"]) | |
| for q in questions] | |
| ms = (time.perf_counter() - t0) * 1000.0 | |
| note = (f" — state truncated from {n_tokens} to {jc.MAX_STATE_TOKENS} tokens (head 60% / tail 40%)" | |
| if truncated else "") | |
| status = ( | |
| f"**{len(questions)} question{'s' if len(questions) != 1 else ''} answered** — " | |
| f"{ms:.0f} ms of GPU compute for {len(questions)} prefill pass" | |
| f"{'es' if len(questions) != 1 else ''}, state {n_tokens} tokens{note}. " | |
| "System 1: one forward pass per question, no tokens generated." | |
| ) | |
| payload = json.dumps( | |
| {"state": state, "system": 1, "answers": answers}, indent=2, ensure_ascii=False | |
| ) | |
| return "".join(_answer_html(q, a) for q, a in zip(questions, answers)), status, payload | |
| def escalate(state: str, question: str, options_text: str, threshold: float, | |
| always_system2: bool, thinking: bool, max_new_tokens: float): | |
| """Confidence-gated escalation: System 1 answers in one pass; System 2 reasons when unsure. | |
| Args: | |
| state: the context to reason about. | |
| question: the question to answer. | |
| options_text: comma-separated options (2-16). Falls back to yes/no if blank. | |
| threshold: escalate to System 2 when System 1's top probability is below this. | |
| always_system2: skip the gate and always ask System 2 to reason. | |
| thinking: run System 2 in thinking mode (step-by-step reasoning). | |
| max_new_tokens: token budget for the System 2 reply. | |
| Returns: | |
| (answer HTML, status, System 2 reply, reasoning, raw JSON) | |
| """ | |
| state = (state or "").strip() | |
| question = (question or "").strip() | |
| if not state: | |
| raise gr.Error("Give the model some state text to reason about.") | |
| if not question: | |
| raise gr.Error("Give the model a question.") | |
| options = _parse_options(options_text) | |
| kind = "choice" if options else "noul" | |
| if kind == "choice" and not 2 <= len(options) <= 16: | |
| raise gr.Error(f"`choice` needs 2-16 comma-separated options; got {len(options)}.") | |
| state_used, _trunc, n_tokens = jc.truncate_state(state) | |
| t0 = time.perf_counter() | |
| probs = jc.decide(kind, state_used, question, options if kind == "choice" else None) | |
| s1_ms = (time.perf_counter() - t0) * 1000.0 | |
| top = max(probs, key=probs.get) | |
| gate = always_system2 or probs[top] < float(threshold) | |
| s1_html = ( | |
| '<div style="padding:10px 0;border-bottom:1px solid var(--border-color-primary)">' | |
| f'<div style="font-size:12.5px;opacity:0.75;margin-bottom:3px">' | |
| f'<code style="font-size:11.5px">system 1</code> {html.escape(question)}</div>' | |
| f'<div style="font-size:19px;font-weight:600;margin:2px 0">{html.escape(str(top))} ' | |
| f'<span style="font-weight:400;font-size:13px;opacity:0.75">' | |
| f"P({html.escape(str(top))}) = {probs[top]:.3f} — one prefill pass</span></div>" | |
| f"{_bars_html(list(probs.items()), top)}</div>" | |
| ) | |
| if not gate: | |
| status = ( | |
| f"**System 1 was confident** — top probability {probs[top]:.3f} ≥ threshold " | |
| f"{float(threshold):.2f}, so System 2 was not consulted. " | |
| f"({s1_ms:.0f} ms of GPU compute, state {n_tokens} tokens.)" | |
| ) | |
| payload = json.dumps({"state": state, "question": question, "options": list(probs), | |
| "system1": probs, "escalated": False}, indent=2, ensure_ascii=False) | |
| return s1_html, status, "", "", payload | |
| opts = list(probs.keys()) | |
| t1 = time.perf_counter() | |
| raw, completed = jc.generate(state_used, question, opts, | |
| max_new_tokens=int(max_new_tokens), | |
| enable_thinking=bool(thinking)) | |
| s2_ms = (time.perf_counter() - t1) * 1000.0 | |
| reasoning, answer = _split_reasoning(raw) | |
| cut_note = "" | |
| if not completed: | |
| cut_note = ( | |
| f"\n\n*System 2 hit its {int(max_new_tokens)}-token budget mid-reply — " | |
| "raise the budget, or untick thinking mode for a concise answer.*" | |
| ) | |
| if thinking and not reasoning: | |
| # the verbose thinking trace never reached its closing tag | |
| reasoning, answer = raw, "" | |
| picked = _matched_option(answer, opts) if completed else None | |
| trigger = ( | |
| f"the 'always ask System 2' box was ticked (System 1's top probability was " | |
| f"{probs[top]:.3f})" | |
| if always_system2 else | |
| f"System 1's top probability {probs[top]:.3f} was below the threshold " | |
| f"{float(threshold):.2f}" | |
| ) | |
| status = ( | |
| f"**Escalated to System 2** — {trigger}. System 1: " | |
| f"{s1_ms:.0f} ms; System 2: {s2_ms:.0f} ms, state {n_tokens} tokens." | |
| + (" Reply cut short at the token budget." if not completed else "") | |
| ) | |
| answer_md = answer | |
| if picked: | |
| answer_md = ( | |
| f"> **System 2 picks: `{picked}`**\n\n" + answer | |
| ) | |
| answer_md += cut_note | |
| payload = json.dumps({"state": state, "question": question, "options": opts, | |
| "system1": probs, "escalated": True, "completed": completed, | |
| "system2_answer": answer, "system2_reasoning": reasoning}, | |
| indent=2, ensure_ascii=False) | |
| return s1_html, status, answer_md, reasoning, payload | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| main, .gradio-container, .gradio-container > .main { max-width: 1150px !important; margin: 0 auto !important; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| # Example rows, taken from the model authors' own published tests: the quickstart cases in | |
| # the model card and the front-page / real-world cases in reports/realworld_9b.json | |
| # (Apache-2.0). Each row is [state, (type, instructions, options) x 4]. | |
| EXAMPLES = [ | |
| [ | |
| "Customer says the parcel arrived damaged and wants their money back.", | |
| "noul", "Is the customer asking for a refund?", "", | |
| "noul", "Is the customer angry or frustrated?", "", | |
| "choice", "Which team should handle this message?", | |
| "billing, logistics & shipping, fraud & abuse, product feedback", | |
| "score", "How polite is the customer, from 0 (abusive) to 5 (very polite)?", "", | |
| ], | |
| [ | |
| "SKU AX-330 stock at 8% of safety level; supplier late twice this quarter.", | |
| "choice", "Supplier response for this scenario.", | |
| "issue_warning, renegotiate, dual_source, maintain", | |
| "noul", "Is the supply risk already material?", "", | |
| "score", "Rate the supply risk on a 0-5 scale.", "", | |
| "noul", "", "", | |
| ], | |
| [ | |
| "Two-tier encryption in the UK", | |
| "choice", "Which topic is this story mainly about?", ", ".join(TOPIC_OPTIONS), | |
| "noul", "Is this story about artificial intelligence or machine learning?", "", | |
| "noul", "", "", | |
| "noul", "", "", | |
| ], | |
| [ | |
| "Please stop this. We've asked you before to observe the guidelines. HN is not a " | |
| "site for political debates, especially ideological ones. We've banned accounts " | |
| "repeatedly for this.", | |
| "choice", "What is this comment mainly doing?", ", ".join(INTENT_OPTIONS), | |
| "noul", "Does this comment insult or personally attack another person or group, " | |
| "rather than just arguing the point?", "", | |
| "noul", "", "", | |
| "noul", "", "", | |
| ], | |
| [ | |
| "STATUS: payments processor is down. Checkout API returns 500s for all " | |
| "transactions; customers cannot pay. Revenue impact growing every minute.", | |
| "score", "Rate operational urgency on a 0-5 scale.", "", | |
| "noul", "Should the on-call engineer be paged now?", "", | |
| "noul", "", "", | |
| "noul", "", "", | |
| ], | |
| ] | |
| with gr.Blocks(title="JEV-9B — typed decisions") as demo: | |
| gr.Markdown( | |
| "# JEV-9B — typed decisions in one forward pass\n" | |
| "An integrated **System 1 + System 2** model built with the Blocks-of-Experts recipe: " | |
| "a frozen **Qwen3.5-9B** backbone, a 40M-parameter LoRA + decision head for " | |
| "**System 1** — typed questions answered in a *single prefill pass* with a **calibrated " | |
| "probability distribution** — and the untouched base model as **System 2** for " | |
| "deliberate step-by-step reasoning. Per the authors' measurements: mean KL ≈ 0.019 to " | |
| "the closed TypeSafe Jev 1.13 on 25k held-out decisions, ECE 0.0007, and HumanEval " | |
| "70.7% — byte-identical to the base model.\n\n" | |
| "Model: [`autotrust/JEV-9B`](https://huggingface.co/autotrust/JEV-9B) (Apache-2.0) · " | |
| "System 2 path = [`Qwen/Qwen3.5-9B`](https://huggingface.co/Qwen/Qwen3.5-9B) unchanged" | |
| ) | |
| with gr.Tabs(): | |
| with gr.TabItem("System 1 — decide"): | |
| gr.Markdown( | |
| "Give the model one piece of **state** text and up to four typed questions. " | |
| "Nothing is generated: the model ranks the options you hand it and returns a " | |
| "probability distribution for each question — one forward pass per question." | |
| ) | |
| state = gr.Textbox( | |
| label="State — the text to reason about", | |
| placeholder="A customer support message, a forum post, an incident summary… " | |
| "up to 1,024 tokens (truncated head 60% / tail 40% beyond that)", | |
| lines=5, max_lines=14, | |
| ) | |
| gr.Markdown( | |
| "#### Questions\n" | |
| "`choice` = one of 2–16 options you list · `score` = a fixed 0–5 ordered " | |
| "scale (describe the scale in the question) · `noul` = a yes/no statement, " | |
| "answered as *P(true)*. Leave a row's instructions blank to skip it." | |
| ) | |
| kind_boxes, instr_boxes, opt_boxes = [], [], [] | |
| for i in range(N_QUESTIONS): | |
| with gr.Row(): | |
| k = gr.Dropdown(choices=KINDS, value="noul", label=f"Q{i + 1} type", | |
| scale=1, allow_custom_value=True, | |
| info="leave the instructions blank to skip a row") | |
| ins = gr.Textbox(label=f"Q{i + 1} instructions", scale=4, lines=1, | |
| placeholder="e.g. Which team should handle this?") | |
| opt = gr.Textbox(label="options (comma-separated, choice only)", scale=5, | |
| lines=1, placeholder="e.g. billing, technical, sales") | |
| kind_boxes.append(k); instr_boxes.append(ins); opt_boxes.append(opt) | |
| run = gr.Button("Decide", variant="primary") | |
| out_status = gr.Markdown() | |
| out_html = gr.HTML(label="Answers") | |
| with gr.Accordion("Raw answers (JSON)", open=False): | |
| out_json = gr.Code(language="json", lines=14, show_label=False, wrap_lines=True) | |
| field_boxes = [c for triple in zip(kind_boxes, instr_boxes, opt_boxes) for c in triple] | |
| decide_inputs = [state] + field_boxes | |
| run.click(decide, inputs=decide_inputs, | |
| outputs=[out_html, out_status, out_json], api_name="decide") | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=decide_inputs, | |
| outputs=[out_html, out_status, out_json], | |
| fn=decide, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples — from the model authors' published tests", | |
| ) | |
| with gr.TabItem("System 1 → System 2 — escalate"): | |
| gr.Markdown( | |
| "The confidence-gated escalation pattern from the model card: System 1 answers " | |
| "in one pass; if its top probability is below your threshold, the *same " | |
| "weights* switch to System 2 — the pristine Qwen3.5-9B with the LoRA off — " | |
| "which thinks it through and writes an answer. Expect the two to disagree " | |
| "sometimes: they are integrated in serving, not in knowledge." | |
| ) | |
| esc_state = gr.Textbox( | |
| label="State — the text to reason about", lines=5, max_lines=14, | |
| placeholder="The context your question is about…", | |
| ) | |
| esc_question = gr.Textbox( | |
| label="Question", lines=1, | |
| placeholder="e.g. Which handler should receive this message?", | |
| ) | |
| esc_options = gr.Textbox( | |
| label="Options (comma-separated, 2–16; leave blank for a yes/no question)", | |
| lines=1, placeholder="e.g. BLOCK, ESCALATE, ROUTE_TECH, ROUTE_SPAM", | |
| ) | |
| with gr.Accordion("Escalation settings", open=True): | |
| esc_threshold = gr.Slider( | |
| 0.5, 1.0, value=0.90, step=0.01, label="Confidence threshold", | |
| info="Escalate to System 2 when System 1's top probability is below this. " | |
| "0.90 is the value suggested in the model card.", | |
| ) | |
| esc_always = gr.Checkbox( | |
| value=False, label="Always ask System 2 (skip the gate)", | |
| ) | |
| esc_thinking = gr.Checkbox( | |
| value=True, label="System 2 thinking mode (step-by-step reasoning)", | |
| ) | |
| esc_tokens = gr.Slider( | |
| 64, 2048, value=1024, step=64, label="System 2 token budget", | |
| info="Thinking mode is verbose — if the reply gets cut short, raise this.", | |
| ) | |
| esc_btn = gr.Button("Ask", variant="primary") | |
| esc_status = gr.Markdown() | |
| esc_html = gr.HTML(label="System 1 — one prefill pass") | |
| esc_answer = gr.Markdown(label="System 2 — deliberate answer") | |
| with gr.Accordion("System 2 reasoning trace", open=False): | |
| esc_reasoning = gr.Markdown(label="Reasoning") | |
| with gr.Accordion("Raw (JSON)", open=False): | |
| esc_json = gr.Code(language="json", lines=12, show_label=False, wrap_lines=True) | |
| esc_btn.click( | |
| escalate, | |
| inputs=[esc_state, esc_question, esc_options, esc_threshold, esc_always, | |
| esc_thinking, esc_tokens], | |
| outputs=[esc_html, esc_status, esc_answer, esc_reasoning, esc_json], | |
| api_name="escalate", | |
| ) | |
| gr.Markdown( | |
| "---\n" | |
| "**What to expect.** System 1 is a distillation of the closed TypeSafe Jev 1.13's " | |
| "published output distributions: it mirrors the teacher's decisions, *including its " | |
| "mistakes*, and inherits its blind spots on arithmetic, dates, counting and " | |
| "adversarial inputs. It reads English best; the first 1,024 tokens of state are " | |
| "used. `choice` handles up to 16 options — beyond ~8, the authors recommend " | |
| "[JEV-27B](https://huggingface.co/autotrust/JEV-27B). Not for high-stakes decisions; " | |
| "use confidence gating.\n\n" | |
| "**Implementation.** Inference follows the authors' documented `transformers` + " | |
| "`peft` path exactly (backbone with the LoRA adapter on, 24-slot fp32 head on the " | |
| "last token, per-kind calibration temperature; System 2 generated with the adapter " | |
| "off). Weights: 18 GB bf16, streamed into VRAM by ZeroGPU on each request. " | |
| "Examples are taken from the model card quickstart and the authors' " | |
| "`reports/realworld_9b.json` (Apache-2.0)." | |
| ) | |
| demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS) |