"""
space/app.py — One for All HuggingFace ZeroGPU Gradio Space.
Run locally (from space/ dir, with a local viz_data.json):
cd space && VIZ_DATA_PATH=/path/to/viz_data.json python app.py
"""
from __future__ import annotations
import os
import html as _html_stdlib
import gradio as gr
import spaces
import _fig
import _glb
import _html
# ── Startup: shared runtime (viz + UMAP + student) lives in _boot ─────────
import _boot
RT = _boot.load_runtime()
HF_TOKEN = RT.hf_token
VIZ = RT.viz
REDUCER = RT.reducer
COORDS3D = RT.coords3d
BACKEND = RT.backend
_MODEL_READY = RT.model_ready
_INIT_GLB = _glb.build_glb(VIZ, COORDS3D, [])
print(f"[ofa-space] GLB path: {_INIT_GLB}")
if _INIT_GLB:
print(f"[ofa-space] GLB exists: {os.path.exists(_INIT_GLB)}, size: {os.path.getsize(_INIT_GLB)} bytes")
# Camera: pull back far enough to see all points at startup
if COORDS3D is not None:
import numpy as _np
_span = float(_np.linalg.norm(COORDS3D.max(axis=0) - COORDS3D.min(axis=0)))
_CAM = (45, 30, _span * 1.8)
else:
_CAM = (45, 30, 10)
def _response_html(text: str, title: str = "MODEL RESPONSE", accent: str = "#8b949e") -> str:
safe = _html_stdlib.escape(text).replace("\n", "
")
return (
'"
)
# ── Backend dispatch: same handler code drives torch and llama.cpp ────────
def _to_device():
_boot.to_device(RT)
def _stream(text: str):
return _boot.stream(RT, text)
def _stream_pair(text: str):
return _boot.stream_pair(RT, text)
def _final_probe(text: str):
return _boot.final_probe(RT, text)
# ── ZeroGPU probe handler (streaming generator) ───────────────────────────
@spaces.GPU
def probe_fn(text: str, probe_points: list):
if not text.strip():
yield gr.skip(), probe_points, "", "", ""
return
if not _MODEL_READY or REDUCER is None:
msg = _html.gate_html([0.2] * 5, VIZ["teacher_names"] or ["—"] * 5)
yield _glb.build_glb(VIZ, COORDS3D, probe_points), probe_points, "", msg, ""
return
_to_device()
names = VIZ["teacher_names"]
# Stream tokens + live gate bars; skip the 3D plot until the end.
partial = ""
for partial, gates in _stream(text):
yield (
gr.skip(), probe_points,
_response_html(partial),
_html.gate_html(gates, names, ranked=False),
gr.skip(),
)
# Final pass: pooled probe point in soul space + dominant-teacher badge.
new_pt, gate_weights = _final_probe(text)
updated = probe_points + [new_pt]
yield (
_glb.build_glb(VIZ, COORDS3D, updated), updated,
_response_html(partial),
_html.gate_html(gate_weights, names),
_html.task_html(gate_weights, names),
)
# ── ZeroGPU arena handler: base (LoRA off) vs deku, same prompt ───────────
@spaces.GPU
def arena_fn(text: str):
if not text.strip():
yield "", "", ""
return
if not _MODEL_READY:
yield "", _response_html("model not loaded", "DEKU · DISTILLED"), ""
return
_to_device()
names = VIZ["teacher_names"]
for base_text, deku_text, gates in _stream_pair(text):
yield (
_response_html(base_text, "BASE · QWEN2.5-0.5B", accent="#8b949e"),
_response_html(deku_text, "DEKU · DISTILLED", accent="#7c3aed"),
_html.gate_html(gates, names, ranked=False),
)
# ── CSS ───────────────────────────────────────────────────────────────────
CSS = """
/* ── Variables ─────────────────────────────────────────── */
:root {
--bg: #080b10;
--panel: #0d1117;
--panel2: #111620;
--border: #1c2129;
--border-hi: #30363d;
--indigo: #7c3aed;
--indigo-dim: rgba(124,58,237,0.18);
--cyan: #06b6d4;
--amber: #f59e0b;
--green: #34d399;
--pink: #f472b6;
--text: #e6edf3;
--text-dim: #8b949e;
--text-faint: #484f58;
--mono: "JetBrains Mono", ui-monospace, SFMono-Regular, monospace;
--radius: 10px;
}
/* ── Base ───────────────────────────────────────────────── */
* { box-sizing: border-box; }
.gradio-container {
background: var(--bg) !important;
background-image:
radial-gradient(ellipse 90% 60% at 10% -5%, rgba(124,58,237,0.07) 0%, transparent 55%),
radial-gradient(ellipse 70% 50% at 90% 105%, rgba(6,182,212,0.05) 0%, transparent 55%)
!important;
font-family: system-ui, -apple-system, BlinkMacSystemFont, sans-serif;
min-height: 100vh;
}
footer { display: none !important; }
/* ── All blocks: remove default boxy look ───────────────── */
.block {
background: var(--panel) !important;
border: 1px solid var(--border) !important;
border-radius: var(--radius) !important;
box-shadow: 0 1px 3px rgba(0,0,0,0.4), 0 0 0 0 transparent !important;
transition: border-color 0.25s, box-shadow 0.25s !important;
padding: 0 !important;
}
.block:hover {
border-color: var(--border-hi) !important;
}
/* Remove the double-border that Gradio adds */
.block .block { border: none !important; background: transparent !important; }
/* ── Label text ─────────────────────────────────────────── */
label > span, .label-wrap span, .block-label span {
font-family: var(--mono) !important;
font-size: 10px !important;
font-weight: 600 !important;
letter-spacing: 0.09em !important;
text-transform: uppercase !important;
color: var(--text-faint) !important;
}
/* ── Textarea / input ───────────────────────────────────── */
textarea, input[type="text"], input[type="number"] {
background: #060a10 !important;
border: 1px solid var(--border-hi) !important;
color: var(--text) !important;
border-radius: 8px !important;
font-size: 13px !important;
line-height: 1.6 !important;
transition: border-color 0.2s, box-shadow 0.2s !important;
resize: vertical !important;
}
textarea:focus, input[type="text"]:focus {
border-color: var(--indigo) !important;
box-shadow: 0 0 0 3px rgba(124,58,237,0.12),
0 0 18px rgba(124,58,237,0.18) !important;
outline: none !important;
}
textarea::placeholder { color: var(--text-faint) !important; }
/* ── Primary button ─────────────────────────────────────── */
button.primary, button[variant="primary"] {
background: linear-gradient(135deg, #6d28d9, var(--indigo)) !important;
border: 1px solid rgba(124,58,237,0.45) !important;
border-radius: 8px !important;
color: #fff !important;
font-family: var(--mono) !important;
font-size: 12px !important;
font-weight: 700 !important;
letter-spacing: 0.07em !important;
text-transform: uppercase !important;
padding: 10px 22px !important;
transition: all 0.2s ease !important;
box-shadow: 0 2px 10px rgba(124,58,237,0.22) !important;
cursor: pointer !important;
}
button.primary:hover {
background: linear-gradient(135deg, var(--indigo), #8b5cf6) !important;
box-shadow: 0 4px 20px rgba(124,58,237,0.45),
0 0 0 1px rgba(124,58,237,0.35) !important;
transform: translateY(-1px) !important;
}
button.primary:active { transform: translateY(0) !important; }
/* Secondary buttons */
button.secondary {
background: var(--panel2) !important;
border: 1px solid var(--border-hi) !important;
color: var(--text-dim) !important;
border-radius: 8px !important;
transition: all 0.2s !important;
}
button.secondary:hover {
border-color: var(--indigo) !important;
color: var(--text) !important;
}
/* ── Tabs: underline style (Linear / GitHub / Vercel) ───── */
.tabs > .tab-nav,
div[role="tablist"] {
background: transparent !important;
border-bottom: 1px solid var(--border) !important;
gap: 0 !important;
padding: 0 2px !important;
}
.tab-nav button, div[role="tab"] {
background: transparent !important;
border: none !important;
border-bottom: 2px solid transparent !important;
border-radius: 0 !important;
color: var(--text-dim) !important;
font-family: var(--mono) !important;
font-size: 11px !important;
font-weight: 600 !important;
letter-spacing: 0.07em !important;
text-transform: uppercase !important;
padding: 10px 18px 9px !important;
margin-bottom: -1px !important;
transition: color 0.18s, border-color 0.18s !important;
box-shadow: none !important;
}
.tab-nav button:hover {
color: var(--text) !important;
border-bottom-color: rgba(124,58,237,0.4) !important;
}
.tab-nav button.selected {
color: var(--indigo) !important;
border-bottom: 2px solid var(--indigo) !important;
background: transparent !important;
box-shadow: none !important;
}
/* ── Plotly / charts: transparent background ────────────── */
.plot-container, .plot-container > div, .js-plotly-plot {
background: transparent !important;
}
/* ── Model3D container ──────────────────────────────────── */
div[data-testid="model3d"], .model3D-component {
border-radius: var(--radius) !important;
overflow: hidden !important;
border: 1px solid var(--border) !important;
box-shadow: 0 0 40px rgba(124,58,237,0.08) inset !important;
}
/* ── Scrollbars ─────────────────────────────────────────── */
::-webkit-scrollbar { width: 5px; height: 5px; }
::-webkit-scrollbar-track { background: var(--bg); }
::-webkit-scrollbar-thumb { background: var(--border-hi); border-radius: 99px; }
::-webkit-scrollbar-thumb:hover { background: var(--text-faint); }
/* ── Animated LIVE badge ─────────────────────────────────── */
@keyframes pulse-dot {
0%, 100% { opacity: 1; box-shadow: 0 0 6px var(--cyan); }
50% { opacity: 0.5; box-shadow: 0 0 2px var(--cyan); }
}
.live-dot { animation: pulse-dot 1.8s ease-in-out infinite; }
/* ── Fade-in on load ─────────────────────────────────────── */
@keyframes fadeUp {
from { opacity: 0; transform: translateY(10px); }
to { opacity: 1; transform: translateY(0); }
}
.gradio-container > .main > .wrap { animation: fadeUp 0.45s ease; }
/* ── Row/col gaps ────────────────────────────────────────── */
.gap { gap: 14px !important; }
.row { gap: 14px !important; }
"""
# ── Layout ────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="One for All") as demo:
gr.HTML(_html.header_html(
n_teachers=len(VIZ["teacher_names"]) or 6,
backend=BACKEND,
))
probe_state = gr.State([])
with gr.Tabs():
# ── Tab 1: Almas ──────────────────────────────────────────────────
with gr.TabItem("Souls"):
with gr.Row():
with gr.Column(scale=6):
umap_plot = gr.Model3D(
value=_INIT_GLB,
display_mode="solid",
clear_color=[0.031, 0.043, 0.063, 1.0],
height=500,
label=None,
camera_position=_CAM,
)
gr.HTML(_glb.build_legend_html(VIZ))
with gr.Column(scale=4):
gr.HTML(
''
'⚡Probe the student'
'LIVE'
'
'
)
prompt_box = gr.Textbox(
lines=4,
placeholder="Ask anything — code, math, language…",
label="",
)
run_btn = gr.Button("Run", variant="primary")
resp_out = gr.HTML()
gate_out = gr.HTML()
task_out = gr.HTML()
gr.HTML(
'↑ new probe point will appear in soul space
'
)
# ── Tab 2: Arena — base vs deku, same prompt ──────────────────────
with gr.TabItem("Arena"):
gr.HTML(
''
'⚔One prompt, two models'
''
'same 0.5B weights — LoRA adapter off vs on'
'
'
)
arena_box = gr.Textbox(
lines=3,
placeholder="Ask something — watch base and distilled race side by side…",
label="",
)
arena_btn = gr.Button("Run both", variant="primary")
with gr.Row():
with gr.Column():
base_out = gr.HTML()
with gr.Column():
deku_out = gr.HTML()
arena_gate_out = gr.HTML()
gr.Examples(
examples=[
"Natalia sold clips to 48 of her friends in April, and then she "
"sold half as many clips in May. How many clips did Natalia sell "
"altogether in April and May?",
"Which property of a mineral can be determined just by looking "
"at it? (A) luster (B) mass (C) weight (D) hardness",
"Write a Python function that checks if a word is a palindrome.",
"Explain why the sky is blue in two sentences.",
],
inputs=[arena_box],
)
# ── Tab 3: Geometria ──────────────────────────────────────────────
with gr.TabItem("Geometry"):
with gr.Row():
with gr.Column(scale=7):
gr.Plot(
value=_fig.build_cka_fig(VIZ["cka"]),
label="CKA geometry alignment",
)
with gr.Column(scale=3):
cka_matrix = VIZ["cka"].get("matrix", [])
if cka_matrix:
import numpy as _np
mat = _np.array(cka_matrix)
n = mat.shape[0]
mask = ~_np.eye(n, dtype=bool)
mean_off = float(mat[mask].mean())
masked = mat.copy()
_np.fill_diagonal(masked, 1.0)
min_idx = _np.unravel_index(masked.argmin(), masked.shape)
hard_pair = (VIZ["cka"]["models"][min_idx[0]],
VIZ["cka"]["models"][min_idx[1]])
hard_val = float(masked[min_idx])
gr.HTML(
f''
f'
{mean_off:.3f}
'
f'
'
f'mean off-diagonal CKA
'
f'
hardest pair
'
f'
'
f'{hard_pair[0]} ↔ {hard_pair[1]}'
f' {hard_val:.2f}
'
f'