ai-model-xray / app.py
CognitiveEngineering's picture
Deploy AI Model X-Ray - structural health scanner
451cf40 verified
Raw
History Blame Contribute Delete
36.6 kB
"""AI Model X-Ray - Structural Health Scanner.
Single-file Gradio app. Upload or select a transformer; get a per-layer
structural-health diagnostic: which layers are compressible (immune), which are
fragile (critical), and how the model compares to reference architectures.
Dark mission-control theme (Octopus palette). All panels are gr.HTML blocks
re-rendered from a scan result produced by analysis.py. Pre-loaded models read
from a precomputed cache (instant, no GPU); custom models trigger a live scan.
Run: python app.py -> http://localhost:7860
"""
from __future__ import annotations
import html
import math
import traceback
import gradio as gr
import analysis as ax
# Optional ZeroGPU decorator (no-op when `spaces` is unavailable, e.g. locally).
try:
import spaces # type: ignore
GPU = spaces.GPU(duration=300)
except Exception: # pragma: no cover - local / non-Spaces environment
def GPU(fn):
return fn
# --------------------------------------------------------------------------- #
# Palette
# --------------------------------------------------------------------------- #
BG = "#0f0f14"
CARD = "#1a1b26"
BORDER = "#2a2b36"
FG = "#f8f8fb"
FG2 = "#8888a0"
EMERALD = "#10B981"
AMBER = "#EF9F27"
RED = "#E24B4A"
REGIME_COLOR = {"immune": EMERALD, "buffer": AMBER, "critical": RED}
DROPDOWN_CHOICES = [
("bert-base-uncased (Encoder, 12L x 12H, 110M)", "bert-base-uncased"),
("gpt2 (Decoder, 12L x 12H, 124M)", "gpt2"),
("google/vit-base-patch16-224 (Vision, 12L x 12H, 86M)", "google/vit-base-patch16-224"),
("distilbert-base-uncased (Encoder, 6L x 12H, 66M)", "distilbert-base-uncased"),
]
# ===========================================================================
# CSS
# ===========================================================================
CSS = """
.gradio-container { background:#0f0f14 !important; color:#f8f8fb !important;
max-width:100% !important;
font-family:system-ui,-apple-system,"Segoe UI",Roboto,sans-serif !important; }
.gradio-container .block, .gradio-container .form, .gradio-container .panel,
.gradio-container .gap { background:transparent !important; border:none !important;
box-shadow:none !important; }
footer { display:none !important; }
.gradio-container button { background:#191a24 !important; color:#cdd2e6 !important;
border:1px solid #2a2b36 !important; box-shadow:none !important;
font-weight:600 !important; transition:border-color .15s ease,color .15s ease; }
.gradio-container button:hover { border-color:#3a3d50 !important; color:#f3f4fa !important; }
/* ---- top bar ---- */
#xr-topbar { display:flex !important; flex-direction:row !important; align-items:center;
gap:10px; padding:12px 18px; width:100%; box-sizing:border-box;
background:linear-gradient(90deg,#15161f,#1a1b26);
border:1px solid #2a2b36; border-radius:14px; margin-bottom:8px; }
#xr-topbar > * { flex-shrink:0; }
#tb-logo-col { flex:1 1 auto; min-width:0; }
.tb-logo { font-size:21px; font-weight:800; color:#f8f8fb; letter-spacing:.02em; }
.tb-logo .xr { color:#10B981; }
.tb-sub { font-size:11.5px; color:#8888a0; margin-top:2px; letter-spacing:.06em;
text-transform:uppercase; }
.tb-badge { display:inline-flex; align-items:center; gap:8px; padding:7px 14px;
border-radius:999px; font-size:12px; font-weight:700; letter-spacing:.04em;
white-space:nowrap; background:rgba(16,185,129,.12); color:#10B981;
border:1px solid #1c5641; }
/* ---- section titles + cards ---- */
.sec-title { font-size:11px; letter-spacing:.13em; text-transform:uppercase;
color:#8888a0; font-weight:800; margin:16px 2px 8px; }
.sec-title .n { color:#10B981; }
.card { background:#1a1b26; border:1px solid #2a2b36; border-radius:14px;
padding:14px 16px; }
/* ---- model info rows ---- */
.kv { display:flex; justify-content:space-between; align-items:center;
padding:7px 2px; border-bottom:1px solid #20212b; font-size:13px; }
.kv:last-child { border-bottom:none; }
.kv .k { color:#8888a0; }
.kv .v { font-weight:800; color:#f8f8fb; font-family:ui-monospace,Consolas,monospace; }
.pill-arch { padding:2px 10px; border-radius:999px; font-size:11px; font-weight:800;
background:rgba(16,185,129,.14); color:#10B981; }
/* ---- arc gauge ---- */
.gauge-wrap { position:relative; width:200px; height:150px; margin:6px auto 0;
animation:pop .5s ease; }
.gauge-svg { width:200px; height:200px; position:absolute; top:0; left:0; }
.g-track { fill:none; stroke:#23242f; stroke-width:13; stroke-linecap:round; }
.g-val { fill:none; stroke-width:13; stroke-linecap:round;
transition:stroke-dashoffset .8s cubic-bezier(.3,1,.4,1),stroke .3s; }
.gauge-center { position:absolute; top:48px; left:0; width:200px; text-align:center; }
.gauge-num { font-size:44px; font-weight:800; line-height:1; }
.gauge-lbl { font-size:12px; letter-spacing:.12em; font-weight:800; margin-top:6px;
text-transform:uppercase; }
.gauge-cap { text-align:center; color:#8888a0; font-size:11px; margin-top:30px;
letter-spacing:.04em; }
@keyframes pop { from{opacity:0;transform:scale(.94)} to{opacity:1;transform:scale(1)} }
/* ---- key metrics ---- */
.metrics { display:grid; grid-template-columns:1fr 1fr; gap:10px; }
.metric { background:#14151d; border:1px solid #2a2b36; border-radius:11px;
padding:11px 12px; }
.metric .mv { font-size:23px; font-weight:800; font-family:ui-monospace,Consolas,monospace;
line-height:1; }
.metric .ml { font-size:10.5px; color:#8888a0; margin-top:6px; letter-spacing:.04em;
text-transform:uppercase; }
/* ---- risk map ---- */
.pills { display:flex; gap:8px; flex-wrap:wrap; margin-bottom:12px; }
.pill { display:inline-flex; align-items:center; gap:7px; padding:6px 13px;
border-radius:999px; font-size:12px; font-weight:700; border:1px solid #2a2b36;
background:#14151d; }
.pill .dot { width:9px; height:9px; border-radius:50%; }
.riskgrid { display:grid; grid-template-columns:repeat(auto-fill,minmax(112px,1fr));
gap:9px; }
.rcell { background:#14151d; border:1px solid #2a2b36; border-radius:11px;
padding:11px 11px 10px; cursor:pointer; position:relative; overflow:hidden;
transition:transform .12s ease,border-color .15s ease; animation:pop .35s ease; }
.rcell:hover { transform:translateY(-2px); border-color:#3a3d50; }
.rcell.sel { border-color:#f8f8fb; box-shadow:0 0 0 2px #f8f8fb inset, 0 6px 18px rgba(0,0,0,.45);
transform:translateY(-2px); }
.rcell.flash { animation:cellflash .45s ease; }
@keyframes cellflash {
0% { box-shadow:0 0 0 2px #f8f8fb inset, 0 0 0 0 rgba(248,248,251,.55); }
100% { box-shadow:0 0 0 2px #f8f8fb inset, 0 0 0 16px rgba(248,248,251,0); } }
.rcell .bar-top { position:absolute; top:0; left:0; right:0; height:3px; }
.rcell .rc-l { font-size:11px; color:#8888a0; font-weight:700; letter-spacing:.04em; }
.rcell .rc-rho { font-size:21px; font-weight:800; margin:5px 0 2px;
font-family:ui-monospace,Consolas,monospace; }
.rcell .rc-badge { font-size:9.5px; font-weight:800; letter-spacing:.05em;
text-transform:uppercase; }
/* ---- spectral profile ---- */
.spec-row { display:flex; align-items:center; gap:9px; margin:6px 0; font-size:12px; }
.spec-lbl { width:34px; color:#8888a0; font-family:ui-monospace,Consolas,monospace;
font-size:11px; flex:0 0 auto; }
.spec-track { flex:1; height:22px; background:#14151d; border-radius:6px; position:relative;
overflow:hidden; }
.spec-fill { height:100%; border-radius:6px; animation:grow .6s ease; transform-origin:left; }
.spec-inval { position:absolute; left:9px; top:0; height:22px; line-height:22px;
font-family:ui-monospace,Consolas,monospace; font-size:11.5px; font-weight:800;
color:#ffffff; text-shadow:0 1px 2px rgba(0,0,0,.6); pointer-events:none; }
@keyframes grow { from{transform:scaleX(.02)} to{transform:scaleX(1)} }
/* color key shown beside the section title */
.spec-key { float:right; font-size:10.5px; font-weight:700; letter-spacing:0;
text-transform:none; color:#cdd2e6; }
.spec-key b { font-weight:700; color:#cdd2e6; }
/* ---- layer detail ---- */
.detail .kv .v.big { font-size:15px; }
.regime-banner { padding:10px 12px; border-radius:11px; font-weight:800; font-size:13px;
margin-bottom:10px; display:flex; align-items:center; justify-content:space-between;
letter-spacing:.03em; }
/* ---- comparison table ---- */
.cmp { width:100%; border-collapse:collapse; font-size:12px; }
.cmp th { text-align:left; color:#8888a0; font-weight:700; padding:6px 7px;
border-bottom:1px solid #2a2b36; font-size:10.5px; letter-spacing:.04em;
text-transform:uppercase; }
.cmp td { padding:7px 7px; border-bottom:1px solid #20212b; color:#dfe2ee;
font-family:ui-monospace,Consolas,monospace; }
.cmp tr.me td { background:rgba(16,185,129,.1); color:#f8f8fb; font-weight:800; }
.cmp tr.me td:first-child { color:#10B981; }
/* ---- pruning rec ---- */
.rec-big { font-size:30px; font-weight:800; color:#10B981; line-height:1;
font-family:ui-monospace,Consolas,monospace; }
.rec-sub { color:#8888a0; font-size:12px; margin:5px 0 12px; }
.rec-list { list-style:none; padding:0; margin:8px 0 0; }
.rec-list li { display:flex; justify-content:space-between; padding:5px 2px;
font-size:12.5px; border-bottom:1px solid #20212b; }
.rec-list li .ln { color:#dfe2ee; }
.rec-list li .hd { color:#8888a0; font-family:ui-monospace,Consolas,monospace; }
.rec-warn { margin-top:12px; padding:10px 12px; border-radius:10px;
background:rgba(226,75,74,.1); border:1px solid rgba(226,75,74,.4);
color:#ff8f8e; font-size:12.5px; font-weight:600; }
.rec-ok { margin-top:12px; padding:10px 12px; border-radius:10px;
background:rgba(16,185,129,.08); border:1px solid rgba(16,185,129,.35);
color:#7fe8c4; font-size:12.5px; font-weight:600; }
/* ---- scan button ---- */
#scan-btn, #scan-btn button { background:linear-gradient(90deg,#0c7a57,#10B981) !important;
color:#fff !important; font-weight:800 !important; letter-spacing:.05em !important;
border:none !important; font-size:14px !important; padding:11px !important; }
#scan-btn:hover, #scan-btn button:hover { filter:brightness(1.09); }
/* dark inputs */
#xr-controls { --input-background-fill:#101019; --block-background-fill:#1a1b26;
--border-color-primary:#2a2b36; --body-text-color:#e8eaf2;
--body-text-color-subdued:#9aa0b5; }
#xr-controls input, #xr-controls textarea, #xr-controls .wrap,
#xr-controls .container { background:#101019 !important; color:#e8eaf2 !important;
border-color:#2a2b36 !important; }
#xr-controls ul.options, #xr-controls .options { background:#101019 !important;
border:1px solid #2a2b36 !important; color:#e8eaf2 !important; }
#xr-controls .options .item:hover, #xr-controls li.item:hover {
background:#20212e !important; }
#xr-controls span, #xr-controls label { color:#c8cce0 !important; }
.muted { color:#8888a0; font-size:13px; padding:14px 4px; text-align:center; }
/* hidden click relay — kept mounted in the DOM (display:none, not unmounted) */
#layer-relay { display:none !important; }
/* ---- footer ---- */
#xr-footer { text-align:center; color:#6f7590; font-size:12px; padding:14px 0 6px;
margin-top:10px; border-top:1px solid #20212b; line-height:1.7; }
#xr-footer b { color:#10B981; }
#xr-footer .law { color:#cdd2e6; font-family:ui-monospace,Consolas,monospace; }
/* ---- top-bar "?" guide button ---- */
#guide-btn, #guide-btn button { width:40px !important; min-width:40px !important;
height:40px !important; border-radius:50% !important; padding:0 !important;
font-size:18px !important; font-weight:800 !important; color:#10B981 !important; }
/* ---- guide modal ---- */
/* No `display` here: Gradio's visible=False sets display:none; when visible the
column is display:flex, so align/justify center the card. (Octopus pattern.) */
#guide-modal { position:fixed; inset:0; z-index:1000; background:rgba(8,8,12,.84);
align-items:center; justify-content:center; padding:30px; }
.modal-card { max-width:700px; max-height:86vh; overflow-y:auto; margin:0 auto;
background:#1a1b26; border:1px solid #2a2b36; border-radius:16px; padding:26px 30px;
box-shadow:0 24px 70px rgba(0,0,0,.6); }
.modal-h1 { font-size:23px; font-weight:800; color:#f8f8fb; margin-bottom:4px; }
.modal-sec { margin-top:16px; }
.modal-h2 { display:inline-block; font-size:12px; font-weight:800; color:#10B981;
letter-spacing:.08em; text-transform:uppercase; margin-bottom:4px; }
.modal-sec p { color:#f8f8fb; font-size:13px; margin:4px 0 0; line-height:1.6; }
.modal-sec ul { margin:7px 0 0; padding-left:18px; color:#f8f8fb; font-size:13px;
line-height:1.75; }
.modal-term { font-weight:800; color:#f8f8fb; font-size:13.5px; margin-top:12px; }
.modal-law { font-family:ui-monospace,Consolas,monospace; color:#cdd2e6; }
.modal-foot { margin-top:20px; padding-top:14px; border-top:1px solid #20212b;
text-align:center; color:#8888a0; font-size:12.5px; line-height:1.7; }
.modal-foot b { color:#10B981; }
.modal-x, .modal-x button { position:fixed !important; top:24px; right:28px; z-index:1001;
width:40px !important; min-width:40px !important; height:40px !important;
flex:none !important; border-radius:10px !important; font-size:16px !important;
background:#1a1b26 !important; border:1px solid #2a2b36 !important;
color:#c8cce0 !important; padding:0 !important; }
.modal-x:hover, .modal-x button:hover { border-color:#E24B4A !important;
color:#ff8f8e !important; }
"""
# JS relay: clicking a risk-map cell writes the layer index into a hidden
# textbox and fires its input event, which drives the Python detail handler.
# Visual feedback (select highlight + flash) is applied immediately on click,
# before the server round-trip, so the response always feels instant.
HEAD_JS = """
<script>
window.xraySelectLayer = function(n, el){
try {
document.querySelectorAll('.rcell.sel').forEach(function(c){ c.classList.remove('sel'); });
if (el) {
el.classList.add('sel');
el.classList.remove('flash'); // restart the animation if re-clicked
void el.offsetWidth;
el.classList.add('flash');
}
} catch (e) {}
var root = document.getElementById('layer-relay');
if (!root) { return; }
var inp = root.querySelector('input, textarea');
if (!inp) { return; }
var proto = inp.tagName === 'TEXTAREA' ? window.HTMLTextAreaElement.prototype
: window.HTMLInputElement.prototype;
var setter = Object.getOwnPropertyDescriptor(proto, 'value').set;
setter.call(inp, String(n) + ':' + Date.now());
inp.dispatchEvent(new Event('input', { bubbles: true }));
};
</script>
"""
# ===========================================================================
# Guide modal content
# ===========================================================================
GUIDE_HTML = """
<div class='modal-card'>
<div class='modal-h1'>How to Read Your Model's X-Ray</div>
<div class='modal-sec'><span class='modal-h2'>What this tool does</span>
<p>AI Model X-Ray analyzes the structural health of transformer models by
examining how attention heads relate to each other. It reveals which
layers are structurally redundant (safe to compress) and which are
critical (removing heads would break the model).</p></div>
<div class='modal-sec'><span class='modal-h2'>The metrics</span>
<div class='modal-term'>Coherence Ratio (&rho;)</div>
<p>Measures how much triangular redundancy exists in each layer's attention
graph. Range 0 to 1.</p>
<ul>
<li><b>&rho; close to 1.0</b>: heads share backup paths. High redundancy.</li>
<li><b>&rho; close to 0.5</b>: moderate redundancy. Prune with caution.</li>
<li><b>&rho; below 0.5</b>: low redundancy. Each head is critical.</li>
<li><b>n/a</b>: not enough triangular structure to measure.</li>
</ul>
<div class='modal-term'>Fragility Index (FI)</div>
<p>Percentage of connections that have NO triangular backup. Lower is
stronger. FI = 0% means every connection has a backup path.</p>
<div class='modal-term'>Structural Integrity</div>
<p>Overall health score (0&ndash;100), computed from the global FI.
100 = fully redundant. Below 50 = structurally fragile.</p>
</div>
<div class='modal-sec'><span class='modal-h2'>The risk map</span>
<p>Each card represents one layer of the model.</p>
<ul>
<li>&#128994; <b>Immune</b> (green): &rho; &gt; 0.8, FI = 0. Safe to prune heads here.</li>
<li>&#128993; <b>Buffer</b> (amber): &rho; 0.5&ndash;0.8. Prune with caution.</li>
<li>&#128308; <b>Critical</b> (red): &rho; &lt; 0.5. Do not remove heads from this layer.</li>
<li>&#11036; <b>n/a</b>: insufficient triangular structure to classify.</li>
</ul>
</div>
<div class='modal-sec'><span class='modal-h2'>The science</span>
<p>Based on the spectral simplicial hierarchy
<span class='modal-law'>&lambda;&#8322;(T(G)) &le; &lambda;&#8322;(G)</span>:
algebraic connectivity can only decrease when moving from the edge graph
to the triangle graph. Validated on 45,000+ graphs with zero violations.
Formally verified in Lean&nbsp;4.</p>
<p>In practice: if a layer's attention heads form many triangles (groups of
3 mutually correlated heads), the layer has structural backup. If they
don't, each head is a single point of failure.</p></div>
<div class='modal-foot'>Built by <b>Cognitive Engineering</b> &#127464;&#127469;<br>
cognitive-engineering.dev</div>
</div>
"""
# ===========================================================================
# Small formatting helpers
# ===========================================================================
def _fnum(x, d=3):
return "&mdash;" if x is None else f"{x:.{d}f}"
def _params_to_millions(params):
"""'110M' -> 110.0, '1.1B' -> 1100.0, None -> None."""
if not params:
return None
s = params.strip().upper()
try:
if s.endswith("B"):
return float(s[:-1]) * 1000.0
if s.endswith("M"):
return float(s[:-1])
return float(s) / 1e6
except ValueError:
return None
# ===========================================================================
# Renderers (pure: result/summary -> HTML)
# ===========================================================================
def render_gauge(score, label, color):
"""270-degree arc gauge for the structural-integrity score."""
r = 80
cx = cy = 100
arc = 0.75 # fraction of the full circle drawn (270 deg)
c = 2 * math.pi * r
track_dash = f"{arc * c:.1f} {c:.1f}"
val_len = arc * c * max(0.0, min(100.0, score)) / 100.0
val_dash = f"{val_len:.1f} {c:.1f}"
# rotate so the 90-degree gap sits centred at the bottom
rot = "transform:rotate(135deg);transform-origin:100px 100px;"
return (
"<div class='gauge-wrap'>"
"<svg class='gauge-svg' viewBox='0 0 200 200'>"
f"<circle class='g-track' cx='{cx}' cy='{cy}' r='{r}' "
f"style='{rot}stroke-dasharray:{track_dash}'></circle>"
f"<circle class='g-val' cx='{cx}' cy='{cy}' r='{r}' "
f"style='{rot}stroke:{color};stroke-dasharray:{val_dash}'></circle>"
"</svg>"
"<div class='gauge-center'>"
f"<div class='gauge-num' style='color:{color}'>{score:.0f}</div>"
f"<div class='gauge-lbl' style='color:{color}'>{label}</div>"
"</div></div>"
"<div class='gauge-cap'>Structural integrity &middot; 100 = fully redundant</div>"
)
def render_model_info(result):
arch = html.escape(result["arch"])
params = result.get("params") or "&mdash;"
src = "cached" if result.get("source") == "cached" else "live scan"
return (
"<div class='card'>"
f"<div class='kv'><span class='k'>Model</span>"
f"<span class='v' style='font-size:12px'>{html.escape(result['model_id'])}</span></div>"
f"<div class='kv'><span class='k'>Parameters</span><span class='v'>{params}</span></div>"
f"<div class='kv'><span class='k'>Layers &times; Heads</span>"
f"<span class='v'>{result['n_layers']} &times; {result['n_heads']}</span></div>"
f"<div class='kv'><span class='k'>Hidden size</span><span class='v'>{result.get('hidden','&mdash;')}</span></div>"
f"<div class='kv'><span class='k'>Architecture</span>"
f"<span class='pill-arch'>{arch}</span></div>"
f"<div class='kv'><span class='k'>Source</span><span class='v' style='font-size:11px'>{src}</span></div>"
"</div>"
)
def render_integrity(summary):
label, color = ax.score_band(summary["score"])
return render_gauge(summary["score"], label, color)
def render_metrics(summary):
rho = summary["rho_mean"]
rho_s = "&mdash;" if rho is None else f"{rho:.2f}"
fi_pct = summary["fi_head"] * 100.0
return (
"<div class='metrics'>"
f"<div class='metric'><div class='mv' style='color:#EF9F27'>{fi_pct:.2f}%</div>"
"<div class='ml'>Fragility index</div></div>"
f"<div class='metric'><div class='mv' style='color:#10B981'>{rho_s}</div>"
"<div class='ml'>Coherence ratio &rho;</div></div>"
f"<div class='metric'><div class='mv'>{summary['n_eligible']}/{summary['n_layers']}</div>"
"<div class='ml'>Eligible layers</div></div>"
f"<div class='metric'><div class='mv' style='color:#10B981'>{summary['violations']}</div>"
"<div class='ml'>Hierarchy violations</div></div>"
"</div>"
)
def render_riskmap(result, summary, selected=None):
pills = (
"<div class='pills'>"
f"<span class='pill'><span class='dot' style='background:#10B981'></span>"
f"{summary['n_immune']} immune</span>"
f"<span class='pill'><span class='dot' style='background:#EF9F27'></span>"
f"{summary['n_buffer']} buffer</span>"
f"<span class='pill'><span class='dot' style='background:#E24B4A'></span>"
f"{summary['n_critical']} critical</span>"
"</div>"
)
cells = []
for rec in result["per_layer"]:
L = rec["layer"]
regime = ax.classify_regime(rec)
color = REGIME_COLOR[regime]
rho = rec.get("rho")
rho_s = "n/a" if rho is None else f"{rho:.2f}"
sel = " sel" if selected == L else ""
cells.append(
f"<div class='rcell{sel}' onclick='xraySelectLayer({L}, this)'>"
f"<div class='bar-top' style='background:{color}'></div>"
f"<div class='rc-l'>LAYER {L}</div>"
f"<div class='rc-rho' style='color:{color}'>{rho_s}</div>"
f"<div class='rc-badge' style='color:{color}'>{regime}</div>"
"</div>"
)
return pills + "<div class='riskgrid'>" + "".join(cells) + "</div>"
def render_spectral(result):
"""Horizontal ρ-per-layer bars. Each bar is colored by regime and scaled
0..1 (so ρ=0.56 is ~56% width); layers where ρ is undefined show a thin gray
'n/a' bar. The ρ value is printed inside the bar, left-aligned."""
rows = []
for rec in result["per_layer"]:
L = rec["layer"]
rho = rec.get("rho")
if rho is None:
color = "#3a3d50" # gray — ρ undefined (T(G) disconnected)
width = 14.0 # thin stub so n/a never reads as a full bar
label = "n/a"
else:
color = REGIME_COLOR[ax.classify_regime(rec)]
width = max(4.0, min(1.0, rho) * 100.0) # 0..1 scale, tiny floor
label = f"{rho:.2f}"
rows.append(
"<div class='spec-row'>"
f"<span class='spec-lbl'>L{L}</span>"
"<div class='spec-track'>"
f"<div class='spec-fill' style='width:{width:.1f}%;background:{color}'></div>"
f"<span class='spec-inval'>{label}</span>"
"</div></div>"
)
return "<div class='card'>" + "".join(rows) + "</div>"
def render_layer_detail(result, selected):
if selected is None or selected >= len(result["per_layer"]):
return ("<div class='card detail'><div class='muted'>Click a layer card in the "
"risk map (or use the selector above) to inspect it.</div></div>")
rec = result["per_layer"][selected]
regime = ax.classify_regime(rec)
meta = ax.REGIME_META[regime]
color = meta["color"]
rho = rec.get("rho")
l2tg = "&mdash;" if not rec["eligible"] else _fnum(rec["l2TG"])
elig = "yes" if rec["eligible"] else "no"
fi_pct = (rec.get("fi") or 0.0) * 100.0
return (
"<div class='card detail'>"
f"<div class='regime-banner' style='background:{color}1f;border:1px solid {color}66;color:{color}'>"
f"<span>Layer {rec['layer']} &middot; {meta['label']}</span>"
f"<span style='font-size:11px;font-weight:700'>{meta['advice']}</span></div>"
f"<div class='kv'><span class='k'>&rho; = &lambda;&#8322;(T)/&lambda;&#8322;(G)</span>"
f"<span class='v big' style='color:{color}'>{'n/a' if rho is None else f'{rho:.3f}'}</span></div>"
f"<div class='kv'><span class='k'>&lambda;&#8322;(G)</span><span class='v'>{_fnum(rec['l2G'])}</span></div>"
f"<div class='kv'><span class='k'>&lambda;&#8322;(T(G))</span><span class='v'>{l2tg}</span></div>"
f"<div class='kv'><span class='k'>Fragility (FI)</span><span class='v'>{fi_pct:.2f}%</span></div>"
f"<div class='kv'><span class='k'>Edges in graph</span><span class='v'>{rec['edges']}</span></div>"
f"<div class='kv'><span class='k'>Heads</span><span class='v'>{result['n_heads']}</span></div>"
f"<div class='kv'><span class='k'>Eligible</span><span class='v'>{elig}</span></div>"
"</div>"
)
def render_comparison(result, summary):
rows = ax.reference_rows()
body = []
for r in rows:
body.append(
"<tr><td>" + html.escape(r["model"]) + "</td>"
f"<td>{_fnum(r['rho_mean'],2)}</td>"
f"<td>{_fnum(r['fi_base'],3)}</td>"
f"<td>{html.escape(r['modality'].title())}</td></tr>"
)
rho = summary["rho_mean"]
me = (
"<tr class='me'><td>&rarr; " + html.escape(result["label"]) + "</td>"
f"<td>{'n/a' if rho is None else f'{rho:.2f}'}</td>"
f"<td>{summary['fi_head']:.3f}</td>"
f"<td>{html.escape(result['modality'].title())}</td></tr>"
)
return (
"<div class='card'><table class='cmp'>"
"<tr><th>Model</th><th>&rho; mean</th><th>FI</th><th>Type</th></tr>"
+ "".join(body) + me + "</table></div>"
)
def render_pruning(result, summary):
immune = [r for r in result["per_layer"] if ax.classify_regime(r) == "immune"]
critical = [r for r in result["per_layer"] if ax.classify_regime(r) == "critical"]
n_heads = result["n_heads"]
total_heads = result["n_layers"] * n_heads
immune_heads = len(immune) * n_heads
pct = (immune_heads / total_heads * 100.0) if total_heads else 0.0
# Rough fp32 memory estimate: immune fraction of total parameters.
pm = _params_to_millions(result.get("params"))
mem = ""
if pm is not None and result["n_layers"]:
frac = len(immune) / result["n_layers"]
saved_mb = pm * 1e6 * frac * 4 / 1e6 # fp32 bytes -> MB
mem = (f"<div class='rec-sub'>Estimated savings: "
f"<b style='color:#f8f8fb'>~{saved_mb:.0f} MB</b> (fp32, "
f"{frac*100:.0f}% of weights)</div>")
items = "".join(
f"<li><span class='ln'>Layer {r['layer']}</span>"
f"<span class='hd'>{n_heads} heads</span></li>" for r in immune
) or "<li><span class='ln' style='color:#8888a0'>No fully-immune layers</span></li>"
if critical:
crit_labels = ", ".join(f"L{r['layer']}" for r in critical)
warn = (f"<div class='rec-warn'>&#9888; Do NOT prune layers: {crit_labels} "
"(low redundancy &mdash; structurally load-bearing).</div>")
else:
warn = ("<div class='rec-ok'>&#10003; No critical layers &mdash; this model has "
"no structurally load-bearing bottlenecks at &rho; &lt; 0.5.</div>")
return (
"<div class='card'>"
f"<div class='rec-big'>{pct:.0f}%</div>"
f"<div class='rec-sub'>of heads ({immune_heads}/{total_heads}) sit in immune "
"layers &mdash; safe to prune.</div>"
+ mem +
"<ul class='rec-list'>" + items + "</ul>"
+ warn +
"</div>"
)
# Empty-state placeholders
EMPTY = "<div class='card'><div class='muted'>Run a scan to populate this panel.</div></div>"
def empty_bundle():
return (EMPTY, render_gauge(0, "&mdash;", FG2), EMPTY, EMPTY, EMPTY, EMPTY, EMPTY, EMPTY)
# ===========================================================================
# Event handlers
# ===========================================================================
# Only the genuine live scan (forward passes) goes through the GPU decorator.
# The cached path — including the BERT auto-load — never touches it, so the app
# works on any hardware (cpu-basic, ZeroGPU) for the pre-loaded models.
@GPU
def _live_scan(model_id):
return ax.scan_model(model_id)
def on_scan(dropdown_id, custom_id, force_rescan):
"""Resolve the model id, scan (cached or live), render every panel."""
model_id = (custom_id or "").strip() or dropdown_id
if not model_id:
gr.Warning("Select a model or enter a HuggingFace model ID.")
return (None, *empty_bundle(), gr.update())
custom = bool((custom_id or "").strip())
want_live = force_rescan or custom
try:
result = None if want_live else ax.cached_result(model_id)
if result is None:
result = _live_scan(model_id)
except Exception as e: # surface a friendly error, keep the app alive
traceback.print_exc()
gr.Warning(f"Scan failed for '{model_id}': {e}")
return (None, *empty_bundle(), gr.update())
summary = ax.summarize(result)
n = result["n_layers"]
layer_choices = [(f"Layer {i}", i) for i in range(n)]
return (
result,
render_model_info(result),
render_integrity(summary),
render_metrics(summary),
render_riskmap(result, summary, None),
render_spectral(result),
render_layer_detail(result, None),
render_comparison(result, summary),
render_pruning(result, summary),
gr.update(choices=layer_choices, value=None),
)
def on_select_dropdown(result, layer_idx):
if result is None or layer_idx is None:
return gr.update(), gr.update()
summary = ax.summarize(result)
return (render_layer_detail(result, int(layer_idx)),
render_riskmap(result, summary, int(layer_idx)))
def on_relay(result, relay_value):
"""Hidden-textbox relay from a risk-map cell click ('<idx>:<ts>')."""
if result is None or not relay_value:
return gr.update(), gr.update(), gr.update()
try:
idx = int(str(relay_value).split(":")[0])
except (ValueError, IndexError):
return gr.update(), gr.update(), gr.update()
summary = ax.summarize(result)
return (render_layer_detail(result, idx),
render_riskmap(result, summary, idx),
gr.update(value=idx))
# ===========================================================================
# App
# ===========================================================================
def build_app():
with gr.Blocks(css=CSS, head=HEAD_JS, title="AI Model X-Ray",
theme=gr.themes.Base()) as demo:
state = gr.State(None)
# ----- top bar ----- #
with gr.Row(elem_id="xr-topbar"):
with gr.Column(elem_id="tb-logo-col"):
gr.HTML("<div class='tb-logo'><span class='xr'>&#128300;</span> "
"AI Model X-Ray</div>"
"<div class='tb-sub'>Structural Health Scanner</div>")
gr.HTML("<span class='tb-badge'>Built by Cognitive Engineering "
"&#127464;&#127469;</span>")
guide_btn = gr.Button("?", elem_id="guide-btn", scale=0, min_width=44)
with gr.Row():
# ===================== LEFT (25%) ===================== #
with gr.Column(scale=25, elem_id="xr-controls"):
gr.HTML("<div class='sec-title'><span class='n'>1.</span> Select model</div>")
dropdown = gr.Dropdown(
choices=DROPDOWN_CHOICES, value="bert-base-uncased",
label="Pre-loaded models", interactive=True)
custom = gr.Textbox(
label="...or a custom HuggingFace model ID", lines=1,
placeholder="e.g. roberta-base, microsoft/deberta-v3-small")
force = gr.Checkbox(label="Force live re-scan (ignore cache)", value=False)
scan_btn = gr.Button("\U0001F52C SCAN MODEL", elem_id="scan-btn")
gr.HTML("<div class='sec-title'>Model info</div>")
info_html = gr.HTML(EMPTY)
gr.HTML("<div class='sec-title'>Structural integrity</div>")
gauge_html = gr.HTML(render_gauge(0, "&mdash;", FG2))
gr.HTML("<div class='sec-title'>Key metrics</div>")
metrics_html = gr.HTML(EMPTY)
# ===================== CENTER (50%) ===================== #
with gr.Column(scale=50):
gr.HTML("<div class='sec-title'><span class='n'>2.</span> Risk map "
"&mdash; per-layer prunability</div>")
riskmap_html = gr.HTML(EMPTY)
gr.HTML("<div class='sec-title'>Spectral profile &mdash; &rho; per layer"
"<span class='spec-key'>&#128994; safe to prune &nbsp; "
"&#128993; caution &nbsp; &#128308; fragile</span></div>")
spectral_html = gr.HTML(EMPTY)
# ===================== RIGHT (25%) ===================== #
with gr.Column(scale=25, elem_id="xr-controls"):
gr.HTML("<div class='sec-title'><span class='n'>3.</span> Layer detail</div>")
layer_sel = gr.Dropdown(choices=[], label="Inspect layer", value=None,
interactive=True)
detail_html = gr.HTML(render_layer_detail({"per_layer": []}, None))
gr.HTML("<div class='sec-title'>Architecture comparison</div>")
comparison_html = gr.HTML(EMPTY)
gr.HTML("<div class='sec-title'>Pruning recommendation</div>")
pruning_html = gr.HTML(EMPTY)
# click-to-select relay for risk-map cells. Kept mounted (CSS display:none
# via #layer-relay) rather than visible=False, so the input element always
# exists in the DOM for the JS to write into.
relay = gr.Textbox(value="", label="", elem_id="layer-relay")
# ----- guide modal (hidden until the "?" button is clicked) ----- #
with gr.Column(visible=False, elem_id="guide-modal") as guide_modal:
guide_close = gr.Button("✕", elem_classes=["modal-x"])
gr.HTML(GUIDE_HTML)
gr.HTML(
"<div id='xr-footer'>"
"<span class='law'>&lambda;&#8322;(T(G)) &le; &lambda;&#8322;(G)</span>: "
"zero violations across 45,000+ graphs and 5 transformer architectures. "
"Formally verified in Lean&nbsp;4.<br>"
"<b>cognitive-engineering.dev</b> &middot; appliedai.ch</div>")
# ----- wiring ----- #
OUTPUTS = [state, info_html, gauge_html, metrics_html, riskmap_html,
spectral_html, detail_html, comparison_html, pruning_html, layer_sel]
scan_btn.click(on_scan, inputs=[dropdown, custom, force], outputs=OUTPUTS)
layer_sel.change(on_select_dropdown, inputs=[state, layer_sel],
outputs=[detail_html, riskmap_html])
relay.input(on_relay, inputs=[state, relay],
outputs=[detail_html, riskmap_html, layer_sel])
# guide modal open/close
guide_btn.click(lambda: gr.update(visible=True), outputs=guide_modal)
guide_close.click(lambda: gr.update(visible=False), outputs=guide_modal)
# Auto-scan the default model (BERT) on page load so the app is never empty.
demo.load(lambda: on_scan("bert-base-uncased", "", False), outputs=OUTPUTS)
return demo
if __name__ == "__main__":
build_app().launch()