File size: 5,540 Bytes
275cef9 956241a 275cef9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | from __future__ import annotations
import numpy as np
import plotly.graph_objects as go
DARK = dict(
paper_bgcolor="#161b22",
plot_bgcolor="#161b22",
font=dict(color="#e6edf3", family="JetBrains Mono, monospace, sans-serif"),
margin=dict(t=48, r=16, b=48, l=56),
)
GRID_COLOR = "#21262d"
# index 0 = student, 1-5 = teachers
MODEL_COLORS = ["#e6edf3", "#7c3aed", "#06b6d4", "#f59e0b", "#34d399", "#f472b6"]
def build_base_traces(viz: dict, coords3d: np.ndarray) -> list:
"""One Scatter3d trace per model. Returns a plain list — never mutate it."""
labels = np.array(viz["labels"])
traces = []
for i, name in enumerate(viz["model_names"]):
c = coords3d[labels == name]
traces.append(go.Scatter3d(
x=c[:, 0].tolist(), y=c[:, 1].tolist(), z=c[:, 2].tolist(),
mode="markers", name=name,
marker=dict(
color=MODEL_COLORS[i % len(MODEL_COLORS)],
size=5 if name != "student" else 6,
opacity=0.85,
),
))
return traces
def rebuild_fig(base_traces: list, probe_points: list[dict]) -> go.Figure:
"""Fresh Figure from immutable base traces + accumulated probe points.
probe_points: list of {"x": float, "y": float, "z": float, "label": str}
Never mutates base_traces.
"""
fig = go.Figure(data=list(base_traces)) # copy, not reference
if probe_points:
fig.add_trace(go.Scatter3d(
x=[p["x"] for p in probe_points],
y=[p["y"] for p in probe_points],
z=[p["z"] for p in probe_points],
mode="markers", name="live probe",
marker=dict(color="#ffffff", size=9, symbol="diamond", opacity=1.0,
line=dict(color="#7c3aed", width=1)),
))
fig.update_layout(
**DARK,
title=dict(text="Soul space — UMAP 3D", font=dict(size=13)),
scene=dict(
bgcolor="#0d1117",
xaxis=dict(showgrid=True, gridcolor=GRID_COLOR, showticklabels=False, title=""),
yaxis=dict(showgrid=True, gridcolor=GRID_COLOR, showticklabels=False, title=""),
zaxis=dict(showgrid=True, gridcolor=GRID_COLOR, showticklabels=False, title=""),
),
legend=dict(bgcolor="rgba(22,27,34,0.85)", bordercolor="#30363d", borderwidth=1,
font=dict(size=11)),
uirevision="soul", # keeps camera position between updates
)
return fig
def build_cka_fig(cka: dict) -> go.Figure:
if not cka or "matrix" not in cka:
return go.Figure(layout={**DARK, "title": "No CKA data"})
fig = go.Figure(go.Heatmap(
z=cka["matrix"], x=cka["models"], y=cka["models"],
colorscale="Viridis", zmin=0, zmax=1,
colorbar=dict(title="CKA", thickness=14, tickfont=dict(color="#e6edf3", size=11)),
text=[[f"{v:.2f}" for v in row] for row in cka["matrix"]],
texttemplate="%{text}", textfont=dict(size=11),
))
fig.update_layout(
**DARK,
title=dict(text="CKA geometry alignment — all pairs", font=dict(size=13)),
xaxis=dict(side="bottom", tickfont=dict(size=11)),
yaxis=dict(autorange="reversed", tickfont=dict(size=11)),
)
return fig
def _ema(vals: list[float], alpha: float = 0.9) -> list[float]:
out, s = [], vals[0]
for v in vals:
s = alpha * s + (1 - alpha) * v
out.append(s)
return out
def build_curves_fig(curves: dict) -> go.Figure:
if not curves or not curves.get("steps"):
return go.Figure(layout={**DARK, "title": "No training data"})
steps = curves["steps"]
fig = go.Figure()
fig.add_trace(go.Scatter(x=steps, y=_ema(curves["task"]), name="task",
line=dict(color="#06b6d4", width=2)))
if curves.get("kd"):
fig.add_trace(go.Scatter(x=steps, y=_ema(curves["kd"]), name="kd (qwen)",
line=dict(color="#34d399", width=2)))
fig.add_trace(go.Scatter(x=steps, y=_ema(curves["geo"]), name="geo",
line=dict(color="#7c3aed", width=2)))
fig.add_trace(go.Scatter(x=steps, y=_ema(curves["total"]), name="total",
line=dict(color="#f59e0b", width=1.5, dash="dot")))
fig.update_layout(
**DARK,
title=dict(text="Loss curves (EMA α=0.9)", font=dict(size=13)),
xaxis=dict(title="step", gridcolor=GRID_COLOR),
yaxis=dict(title="loss", gridcolor=GRID_COLOR),
legend=dict(bgcolor="rgba(22,27,34,0.8)", bordercolor="#30363d", borderwidth=1),
)
return fig
def build_gate_area_fig(curves: dict) -> go.Figure:
if not curves or not curves.get("gate"):
return go.Figure(layout={**DARK, "title": "No gate data"})
steps = curves["steps"]
names = curves["teacher_names"]
teacher_colors = MODEL_COLORS[1:]
fig = go.Figure()
for i, name in enumerate(names):
fig.add_trace(go.Scatter(
x=steps,
y=[g[i] for g in curves["gate"]],
name=name, mode="lines", stackgroup="g",
line=dict(color=teacher_colors[i % len(teacher_colors)], width=0.5),
))
fig.update_layout(
**DARK,
title=dict(text="Gate — teacher routing over steps", font=dict(size=13)),
yaxis=dict(title="weight", range=[0, 1], gridcolor=GRID_COLOR),
xaxis=dict(title="step", gridcolor=GRID_COLOR),
legend=dict(bgcolor="rgba(22,27,34,0.8)", bordercolor="#30363d", borderwidth=1),
)
return fig
|