"""
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 _data
import _fig
import _glb
import _html
import _probe
# ── Startup: load data, fit UMAP, load student ───────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN")
_local_viz = os.environ.get("VIZ_DATA_PATH")
try:
if _local_viz:
VIZ = _data.load_from_path(_local_viz)
print(f"[ofa-space] loaded viz from {_local_viz}")
else:
VIZ = _data.load_and_parse(HF_TOKEN)
except Exception as e:
print(f"[ofa-space] viz_data.json not available ({e}), using empty state")
VIZ = _data.make_empty_viz()
try:
if VIZ["stacked"].shape[0] > 3:
REDUCER = _data.fit_umap3d(VIZ["stacked"])
COORDS3D = REDUCER.embedding_
print(f"[ofa-space] UMAP done: {COORDS3D.shape}")
else:
print(f"[ofa-space] not enough points for UMAP: {VIZ['stacked'].shape[0]}")
REDUCER = None
COORDS3D = None
except Exception as e:
print(f"[ofa-space] UMAP failed ({e}), 3D disabled")
REDUCER = None
COORDS3D = None
try:
TOK, STUDENT, GATING = _probe.load_student(HF_TOKEN)
_MODEL_READY = True
except Exception as e:
print(f"[ofa-space] Student not available ({e}). Probe disabled.")
TOK = STUDENT = GATING = None
_MODEL_READY = False
_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) -> str:
safe = _html_stdlib.escape(text).replace("\n", "
")
return (
''
'
MODEL RESPONSE
'
f'
{safe}
'
"
'
'⚡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: 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'