app.py

"""QRoute HuggingFace Space - quantum MoE router demo, runs on CPU Basic.

The Space serves *precomputed* toy-MoE results (both routers trained to the README
config, baked by scripts/precompute_space.py) and renders them with matplotlib. It
does no live training or inference, so torch / pennylane / qroute are not in the
image. There is no IBM path and no OAuth: the Space is fully open.

Design (see README.md): a hero with the 3-qubit VQC router, an interactive
expert-routing explorer (8 expert "lamps" lit by which router selects them top-2,
plus side-by-side routing-probability bars), training loss curves, expert
utilization with Gini load-balance, the toy-MoE results table, and an honest
"feasibility, not advantage" framing.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import gradio as gr
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = Path(__file__).resolve().parent
RESULTS_PATH = HERE / "precomputed" / "toy_results.json"
ASSETS = HERE / "assets"
REPO_URL = "https://github.com/Quantum-Labor/qroute"

CYAN, AMBER, VIOLET, GREEN, MUTED, TEXT = (
    "#22D3EE",
    "#F59E0B",
    "#7C3AED",
    "#34D399",
    "#9CA3AF",
    "#E5E7EB",
)


def _load_results() -> dict[str, Any]:
    try:
        return json.loads(RESULTS_PATH.read_text(encoding="utf-8"))  # type: ignore[no-any-return]
    except (FileNotFoundError, json.JSONDecodeError):
        return {"config": {}, "quantum": {}, "classical": {}, "examples": []}


RESULTS = _load_results()
EXAMPLES: list[dict[str, Any]] = RESULTS.get("examples", [])
N_EXPERTS: int = RESULTS.get("config", {}).get("n_experts", 8)
EX_BY_ID: dict[str, dict[str, Any]] = {e["id"]: e for e in EXAMPLES}


def _read_asset(name: str) -> str:
    try:
        return (ASSETS / name).read_text(encoding="utf-8")
    except FileNotFoundError:
        return ""


# --- rendering helpers (pure, unit-testable) --------------------------------


def example_choices() -> list[str]:
    return [e["id"] for e in EXAMPLES]


def render_lamps(example: dict[str, Any]) -> str:
    vqc = set(example["quantum"]["top2"])
    cls = set(example["classical"]["top2"])
    qp = example["quantum"]["probs"]
    cp = example["classical"]["probs"]
    lamps: list[str] = []
    for e in range(N_EXPERTS):
        klass = "qr-lamp"
        if e in vqc and e in cls:
            klass += " both"
        elif e in vqc:
            klass += " vqc"
        elif e in cls:
            klass += " classical"
        pct = f"{qp[e] * 100:.0f} / {cp[e] * 100:.0f}%"
        lamps.append(
            f'<div class="{klass}"><div class="bulb"></div>'
            f'<div class="name">e{e}</div><div class="pct">{pct}</div></div>'
        )
    legend = (
        '<div class="qr-legend">'
        '<span><span class="dot" style="background:#22D3EE"></span>VQC top-2</span>'
        '<span><span class="dot" style="background:#F59E0B"></span>classical top-2</span>'
        '<span><span class="dot" style="background:linear-gradient(135deg,#22D3EE 50%,#F59E0B 50%)">'  # noqa: E501
        "</span>both agree</span></div>"
    )
    return f'<div class="qr-lamps">{"".join(lamps)}</div>{legend}'


def render_summary(example: dict[str, Any]) -> str:
    vqc = example["quantum"]["top2"]
    cls = example["classical"]["top2"]
    agree = sorted(set(vqc) & set(cls))
    return (
        '<div class="qr-scores">'
        f'<div class="qr-score-card vqc"><h4>VQC router top-2</h4>'
        f'<div class="qr-score-val">{vqc}</div>'
        f'<div class="qr-score-sub">3-qubit circuit, 8 basis-state probabilities</div></div>'
        f'<div class="qr-score-card classical"><h4>Classical router top-2</h4>'
        f'<div class="qr-score-val">{cls}</div>'
        f'<div class="qr-score-sub">MLP baseline · agreement: {agree or "none"}</div></div>'
        "</div>"
    )


def _dark_ax(ax: Any, fig: Any, title: str) -> None:
    fig.patch.set_facecolor("#0B0B16")
    ax.set_facecolor("#0B0B16")
    ax.set_title(title, color=TEXT, fontsize=12, pad=10)
    ax.tick_params(colors="#6B6788", labelsize=8)
    for spine in ax.spines.values():
        spine.set_color("#2A2440")


def routing_figure(example: dict[str, Any]) -> Any:
    plt.style.use("dark_background")
    fig, ax = plt.subplots(figsize=(7.6, 3.0), dpi=110)
    idx = list(range(N_EXPERTS))
    qp = example["quantum"]["probs"]
    cp = example["classical"]["probs"]
    width = 0.4
    ax.bar([i - width / 2 for i in idx], qp, width, color=CYAN, label="VQC")
    ax.bar([i + width / 2 for i in idx], cp, width, color=AMBER, label="classical")
    ax.set_xticks(idx)
    ax.set_xticklabels([f"e{i}" for i in idx], fontsize=8)
    ax.set_ylabel("routing probability", color=MUTED, fontsize=10)
    _dark_ax(ax, fig, "Routing probabilities: VQC vs classical")
    ax.legend(facecolor="#15152A", edgecolor="#2A2440", labelcolor=TEXT, fontsize=9)
    fig.tight_layout()
    return fig


def loss_figure() -> Any:
    plt.style.use("dark_background")
    fig, ax = plt.subplots(figsize=(7.6, 3.0), dpi=110)
    q = RESULTS.get("quantum", {}).get("losses", [])
    c = RESULTS.get("classical", {}).get("losses", [])
    if q:
        ax.plot(range(len(q)), q, color=CYAN, linewidth=2, label="VQC")
    if c:
        ax.plot(range(len(c)), c, color=AMBER, linewidth=2, label="classical")
    ax.set_xlabel("epoch", color=MUTED, fontsize=10)
    ax.set_ylabel("training loss", color=MUTED, fontsize=10)
    _dark_ax(ax, fig, "Training loss: both routers converge")
    if q or c:
        ax.legend(facecolor="#15152A", edgecolor="#2A2440", labelcolor=TEXT, fontsize=9)
    fig.tight_layout()
    return fig


def utilization_figure() -> Any:
    plt.style.use("dark_background")
    fig, ax = plt.subplots(figsize=(7.6, 3.0), dpi=110)
    q = RESULTS.get("quantum", {})
    c = RESULTS.get("classical", {})
    qu = q.get("utilization", [0] * N_EXPERTS)
    cu = c.get("utilization", [0] * N_EXPERTS)
    idx = list(range(N_EXPERTS))
    width = 0.4
    ax.bar(
        [i - width / 2 for i in idx], qu, width, color=CYAN, label=f"VQC (Gini {q.get('gini', 0)})"
    )
    ax.bar(
        [i + width / 2 for i in idx],
        cu,
        width,
        color=AMBER,
        label=f"classical (Gini {c.get('gini', 0)})",
    )
    ax.set_xticks(idx)
    ax.set_xticklabels([f"e{i}" for i in idx], fontsize=8)
    ax.set_ylabel("times selected (val set)", color=MUTED, fontsize=10)
    _dark_ax(ax, fig, "Expert utilization (lower Gini = more balanced)")
    ax.legend(facecolor="#15152A", edgecolor="#2A2440", labelcolor=TEXT, fontsize=9)
    fig.tight_layout()
    return fig


def select_example(example_id: str) -> tuple[str, str, Any]:
    example = EX_BY_ID[example_id]
    return render_lamps(example), render_summary(example), routing_figure(example)


def results_markdown() -> str:
    q = RESULTS.get("quantum", {})
    c = RESULTS.get("classical", {})
    cfg = RESULTS.get("config", {})
    if not q:
        return "_Results unavailable._"

    def row(name: str, r: dict[str, Any]) -> str:
        cells = f"{r.get('params')} | {r.get('val_acc')} | {r.get('train_acc')} | {r.get('gini')}"
        return f"| {name} | {cells} |"

    return (
        "| Router | Params | Val acc | Train acc | Expert Gini |\n"
        "| --- | --- | --- | --- | --- |\n"
        + row("VQC (3-qubit)", q)
        + "\n"
        + row("classical (MLP)", c)
        + "\n\n"
        + f"8 experts, top-2 routing, {cfg.get('n_clusters')}-cluster synthetic task, "
        f"{cfg.get('epochs')} epochs, seed {cfg.get('seed')}. **Feasibility, not advantage:** "
        "both routers converge to the same accuracy at a comparable parameter count; the "
        "VQC router is a trainable, integrable quantum module, not a faster or more accurate "
        "one. On this toy run the VQC happens to balance expert load more evenly (lower Gini), "
        "which is an observation, not a claim."
    )


# --- UI ---------------------------------------------------------------------

_THEME = gr.themes.Base(
    primary_hue="purple",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
    body_background_fill="#0B0B16",
    body_text_color="#E5E7EB",
    background_fill_primary="#15152A",
    background_fill_secondary="#0B0B16",
    border_color_primary="#2A2440",
    button_primary_background_fill="#7C3AED",
    button_primary_background_fill_hover="#8B5CF6",
    button_primary_text_color="#FFFFFF",
    block_background_fill="#15152A",
    block_border_color="#2A2440",
    input_background_fill="#15152A",
)


def _hero_html() -> str:
    hero = _read_asset("hero_vqc.svg")
    art = f'<div class="qr-hero-art">{hero}</div>' if hero else '<div class="qr-hero-art"></div>'
    return f"""
<div class="qr-hero"><div class="qr-hero-inner">
  {art}
  <div class="qr-hero-copy">
    <h1 class="qr-title">QRoute</h1>
    <p class="qr-tagline">A variational quantum circuit that routes tokens to experts
    in a Mixture-of-Experts LLM.</p>
    <p class="qr-sub">n qubits give 2^n basis states - one per expert - so a tiny
    circuit scores an exponential number of experts. This Phase-1 prototype routes
    top-2 of 8 experts with a 3-qubit VQC, alongside a classical baseline.</p>
    <div class="qr-badges">
      <span class="qr-badge">Phase 1</span>
      <a class="qr-badge" href="{REPO_URL}">GitHub</a>
      <span class="qr-badge cyan">Project 3 of 3</span>
      <span class="qr-badge green">simulator</span>
      <span class="qr-badge amber">Apache-2.0</span>
    </div>
  </div>
</div></div>
"""


_WHAT = """
### What is this?

A Mixture-of-Experts (MoE) layer routes each token to a few of many expert
networks. The router is a small classifier - and a natural place to ask whether a
quantum circuit can do the job, because `n` qubits give `2^n` computational basis
states, one per expert, so a small circuit can score an exponential number of
experts. QRoute prototypes this on a toy MoE: a 3-qubit variational quantum
circuit routes the top-2 of 8 experts, with a classical MLP router as a baseline.
"""

_HOW = """
### How it works (in 30 seconds)

- A learned linear layer maps a token embedding to `2 x n_qubits` rotation angles.
- A VQC (data encoding + variational layers + a CNOT ring) is measured with
  `qml.probs`, giving `2^n_qubits` expert routing scores from `n_qubits` qubits.
- Gumbel-softmax top-k selects the experts (differentiable in training, hard at
  inference).
- The MoE output is the routing-weighted sum of the selected experts.

This Space serves results precomputed with that pipeline; nothing trains live.
"""

_ABOUT = f"""
### About

QRoute is project 3 of 3 in the Quantum Co-Processor program, after
[QVerify](https://github.com/Quantum-Labor/qverify) (Grover-assisted reasoning
verification) and [QAgent](https://github.com/Quantum-Labor/qagent) (QAOA tool
selection). Source, the full-scale 128-expert design, and the roadmap:
[{REPO_URL}]({REPO_URL}).

**Honest scope.** Simulator only, 3 qubits / 8 experts, a toy task; no quantum
advantage is claimed. The value of Phase 1 is a trainable, integrable quantum
router and an honest baseline. Scaling to a 7-qubit / 128-expert router in Gemma 4
is the planned Phase 2-4 work (see the repo's design.md and roadmap.md).
"""


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="QRoute - quantum MoE router") as demo:
        gr.HTML(_hero_html())
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(_WHAT)
            with gr.Column(scale=1):
                gr.Markdown(_HOW)

        gr.HTML('<div class="qr-section-title">Try it now: expert routing</div>')
        default_id = example_choices()[0] if EXAMPLES else None
        ex_dd = gr.Dropdown(
            choices=example_choices(),
            value=default_id,
            label="Example token",
            info="Each token comes from a different cluster; see which experts each router lights.",
        )
        lamps = gr.HTML()
        summary = gr.HTML()
        routing = gr.Plot(label="Routing probabilities")

        gr.HTML('<div class="qr-section-title">Training and utilization</div>')
        with gr.Row():
            gr.Plot(value=loss_figure(), label="Training loss")
            gr.Plot(value=utilization_figure(), label="Expert utilization")

        gr.HTML('<div class="qr-section-title">Toy MoE results</div>')
        gr.Markdown(results_markdown())

        gr.Markdown(_ABOUT)

        ex_dd.change(select_example, inputs=[ex_dd], outputs=[lamps, summary, routing])
        if default_id is not None:
            demo.load(select_example, inputs=[ex_dd], outputs=[lamps, summary, routing])
    return demo


demo = build_demo()

if __name__ == "__main__":
    demo.launch(
        theme=_THEME,
        css=_read_asset("styles.css"),
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
    )