"""
app.py
Gradio Space: Interactive Transformer Visualizer — English → Bengali
"""
import gradio as gr
import torch
import json
import os
import numpy as np
from pathlib import Path
from transformer import Transformer, CalcLog
from training import build_model, run_training, visualize_training_step, collate_batch
from inference import visualize_inference
from vocab import get_vocabs, PARALLEL_DATA, PAD_IDX
# ─────────────────────────────────────────────
# Global state
# ─────────────────────────────────────────────
DEVICE = "cpu"
src_v, tgt_v = get_vocabs()
MODEL: Transformer = None
LOSS_HISTORY = []
IS_TRAINED = False
def get_or_init_model():
global MODEL
if MODEL is None:
MODEL = build_model(len(src_v), len(tgt_v), DEVICE)
return MODEL
# ─────────────────────────────────────────────
# HTML renderer for calc log
# ─────────────────────────────────────────────
def render_matrix_html(val, max_rows=6, max_cols=8):
"""Convert a nested list / scalar to an HTML matrix table."""
if isinstance(val, (int, float)):
return f'{val:.5f}'
if isinstance(val, dict):
rows = "".join(
f'
{k}
{v}
'
for k, v in val.items()
)
return f'
{rows}
'
if isinstance(val, list):
# 0-D or scalar list
if len(val) == 0:
return "empty"
# 1-D
if not isinstance(val[0], list):
clipped = val[:max_cols*2]
cells = "".join(
f'
{v:.4f}
'
if isinstance(v, float) else f'
{v}
'
for v in clipped
)
suffix = f'
…+{len(val)-len(clipped)}
' if len(val) > len(clipped) else ""
return f'
{cells}{suffix}
'
# 2-D
rows_html = ""
display_rows = val[:max_rows]
for row in display_rows:
display_cols = row[:max_cols]
cells = "".join(
f'
'
f'{float(c):.3f}
'
if isinstance(c, (int, float)) else f'
{c}
'
for c in display_cols
)
suffix = f'
…
' if len(row) > max_cols else ""
rows_html += f"
{cells}{suffix}
"
if len(val) > max_rows:
rows_html += f'
…{len(val)-max_rows} more rows
'
return f'
{rows_html}
'
return f'{str(val)[:200]}'
def calc_log_to_html(steps):
"""Turn CalcLog steps into rich HTML accordion."""
if not steps:
return "
No calculation log yet.
"
cards = []
for i, step in enumerate(steps):
name = step.get("name", f"step_{i}")
formula = step.get("formula", "")
note = step.get("note", "")
shape = step.get("shape")
val = step.get("value")
shape_badge = f'{shape}' if shape else ""
formula_html = f'
⟨ {formula} ⟩
' if formula else ""
note_html = f'
ℹ {note}
' if note else ""
matrix_html = render_matrix_html(val) if val is not None else ""
# Color category by name prefix
cat = "default"
n = name.upper()
if "EMBED" in n or "TOKEN" in n: cat = "embed"
elif "PE" in n or "POSITIONAL" in n: cat = "pe"
elif "SOFTMAX" in n or "ATTN" in n or "_Q" in n or "_K" in n or "_V" in n: cat = "attn"
elif "FFN" in n or "LINEAR" in n or "RELU" in n: cat = "ffn"
elif "NORM" in n or "RESIDUAL" in n: cat = "norm"
elif "LOSS" in n or "GRAD" in n or "OPTIM" in n: cat = "loss"
elif "INFERENCE" in n or "GREEDY" in n or "BEAM" in n: cat = "infer"
elif "CROSS" in n: cat = "cross"
elif "MASK" in n: cat = "mask"
cards.append(f"""
#{i+1}{name.replace('_',' ')}
{shape_badge}
▶
{formula_html}
{note_html}
{matrix_html}
""")
return "\n".join(cards)
# ─────────────────────────────────────────────
# Attention heatmap HTML
# ─────────────────────────────────────────────
def attention_heatmap_html(weights, row_labels, col_labels, title="Attention"):
"""weights: 2D list [tgt, src]"""
if not weights:
return ""
rows_html = ""
for i, row in enumerate(weights):
cells = ""
for j, w in enumerate(row):
alpha = min(float(w), 1.0)
cells += f'
{alpha:.2f}
'
lbl = row_labels[i] if i < len(row_labels) else str(i)
rows_html += f'
{lbl}
{cells}
'
header = '
' + "".join(f'
{c}
' for c in col_labels) + '
'
return f"""
{title}
{header}{rows_html}
"""
# ─────────────────────────────────────────────
# Decoding steps HTML
# ─────────────────────────────────────────────
def decode_steps_html(step_logs, src_tokens):
if not step_logs:
return ""
html = '
🔁 Auto-regressive Decoding Steps
'
for s in step_logs:
step = s.get("step", 0)
tokens_so_far = s.get("tokens_so_far", [])
top5 = s.get("top5", [])
chosen = s.get("chosen_token", "?")
prob = s.get("chosen_prob", 0)
bars = ""
if top5:
max_p = max(t["prob"] for t in top5) or 1
for t in top5:
pct = t["prob"] / max_p * 100
is_chosen = "chosen" if t["token"] == chosen else ""
bars += f"""
{t['token']}{t['prob']:.3f}
"""
cross_heat = ""
if s.get("cross_attn") and src_tokens:
attn_mat = s["cross_attn"] # [num_heads][T_q][T_src]
if attn_mat and attn_mat[0]:
# Take head-0, last decoded position → [T_src] floats
last_pos_attn = attn_mat[0][-1] # [T_src]
last_row = [last_pos_attn] # [[T_src]] — 2D for heatmap
cross_heat = attention_heatmap_html(
last_row, [chosen], src_tokens,
title=f"Cross-Attn: '{chosen}' → English"
)
html += f"""