#!/usr/bin/env python3
"""
visualizar_flujo.py — Visualizador interactivo del flujo de datos en PamparV3.
Genera un HTML standalone con 4 paneles:
1. Embeddings iniciales (PCA 2D) — dónde vive cada token en el espacio
2. Mapa de atención por nivel — qué tokens miran a cuáles
3. Norma de activaciones por nivel/stream — cómo crece/decae la señal
4. Benchmark de velocidad y memoria
Uso:
python visualizar_flujo.py "Hola, me llamo Pampar"
python visualizar_flujo.py --texto "def suma(a, b): return a + b"
python visualizar_flujo.py # usa texto por defecto
"""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
import torch
import torch.nn.functional as F
# ── Setup de paths ──────────────────────────────────────────────────────────
ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(Path(__file__).parent))
# ── Captura de activaciones via hooks ───────────────────────────────────────
class ActivationCapture:
"""Registra hooks en PamparV3 y captura tensores del forward pass."""
def __init__(self):
self.handles: list = []
self.embeddings: torch.Tensor | None = None # [B, L, dim]
self.nivel_outputs: list[torch.Tensor] = [] # por nivel: [B, L, dim]
self.attn_weights: list[torch.Tensor] = [] # por nivel: [B, H, L, L]
self.stream_norms: list[list[float]] = [] # por nivel: [4 streams]
def attach(self, model) -> None:
"""Registra hooks en el modelo."""
# Hook en embedding
def hook_emb(module, inp, out):
self.embeddings = out.detach().cpu()
self.handles.append(model.tok_emb.register_forward_hook(hook_emb))
# Hook en cada NivelProfundo
for i, nivel in enumerate(model.niveles):
def make_nivel_hook(idx):
def hook(module, inp, out):
# out = (streams_list, terr_acts, conf)
# streams_list es una lista de n_streams tensores [B, L, D]
streams_out = out[0] if isinstance(out, (tuple, list)) else out
if isinstance(streams_out, (list, tuple)):
x = torch.stack([s.detach().cpu() for s in streams_out]).mean(0)
else:
x = streams_out.detach().cpu()
self.nivel_outputs.append(x)
# Norma por stream (cada cuarto del dim como proxy de stream)
B, L, D = x.shape
chunk = max(1, D // 4)
norms = [
x[:, :, i * chunk : min((i + 1) * chunk, D)]
.norm(dim=-1)
.mean()
.item()
for i in range(4)
]
self.stream_norms.append(norms)
return hook
self.handles.append(nivel.register_forward_hook(make_nivel_hook(i)))
# Hook en la atención del nivel
def make_attn_hook(idx):
def hook(module, inp, out):
# Recalcular pesos de atención desde el input
x = inp[0]
B, L, D = x.shape
H = module.n_heads
Hkv = module.n_kv_heads
head_dim = module.head_dim
# Q, K usando nombres reales del BloqueAttn
q = module.q_proj(x).view(B, L, H, head_dim).transpose(1, 2)
k = module.k_proj(x).view(B, L, Hkv, head_dim).transpose(1, 2)
k = module._repeat_kv(k) # GQA expand
scale = head_dim**-0.5
scores = (
torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale
)
# Máscara causal
mask = torch.triu(
torch.ones(L, L, device=x.device), diagonal=1
).bool()
scores = scores.masked_fill(
mask.unsqueeze(0).unsqueeze(0), float("-inf")
)
weights = F.softmax(scores, dim=-1)
self.attn_weights.append(weights.detach().cpu())
return hook
self.handles.append(nivel.attn.register_forward_hook(make_attn_hook(i)))
def detach(self) -> None:
for h in self.handles:
h.remove()
self.handles.clear()
# ── PCA manual (sin sklearn) ────────────────────────────────────────────────
def pca_2d(matrix: torch.Tensor) -> torch.Tensor:
"""Reduce [N, D] a [N, 2] via PCA (SVD)."""
m = matrix.float()
m = m - m.mean(0, keepdim=True)
_, _, V = torch.svd(m)
return m @ V[:, :2]
# ── Generación del HTML ─────────────────────────────────────────────────────
def build_html(
tokens: list[str],
capture: ActivationCapture,
text: str,
elapsed_ms: float,
mem_mb: float,
) -> str:
n_tokens = len(tokens)
n_levels = len(capture.attn_weights)
# ── Datos embedding PCA ─────────────────────────────────────────────────
emb = capture.embeddings[0] # [L, dim]
if n_tokens >= 2:
coords = pca_2d(emb).tolist()
else:
coords = [[0.0, 0.0]] * n_tokens
emb_x = [c[0] for c in coords]
emb_y = [c[1] for c in coords]
emb_norm = emb.norm(dim=-1).tolist()
# ── Datos atención — promedio de cabezas por nivel ───────────────────────
attn_data = []
for lvl_w in capture.attn_weights:
# lvl_w: [B, H, L, L] → promedio de H → [L, L]
avg = lvl_w[0].mean(0).tolist()
attn_data.append(avg)
# ── Normas por nivel/stream ──────────────────────────────────────────────
stream_norms = capture.stream_norms # [[s0,s1,s2,s3], ...] por nivel
# ── Serializar a JSON inline ─────────────────────────────────────────────
import json
tokens_json = json.dumps(tokens)
emb_x_json = json.dumps(emb_x)
emb_y_json = json.dumps(emb_y)
emb_norm_json = json.dumps(emb_norm)
attn_json = json.dumps(attn_data)
stream_norms_json = json.dumps(stream_norms)
n_levels_json = n_levels
elapsed_json = elapsed_ms
mem_json = mem_mb
text_escaped = text.replace('"', '\\"')
return f"""
PamparV3 — Flujo de datos
PamparV3 — Visualizador de Flujo Interno
Arquitectura: 640d · {n_levels_json} niveles · 4 streams · GQA
"{text_escaped}"
{elapsed_ms:.1f}ms
forward pass
{mem_mb:.0f}MB
VRAM usada
{n_levels_json}
niveles de profundidad
1. Embeddings iniciales (PCA 2D)
3. Norma de activaciones por stream
4. Distancia que recorre cada token
"""
# ── Main ─────────────────────────────────────────────────────────────────────
def main():
parser = argparse.ArgumentParser(description="Visualizador de flujo PamparV3")
parser.add_argument(
"texto",
nargs="?",
default="Hola me llamo Pampar y aprendo a programar",
help="Texto a analizar",
)
parser.add_argument(
"--texto", dest="texto_flag", help="Alternativa: --texto 'tu frase'"
)
parser.add_argument("--checkpoint", default="checkpoints/v3_classroom.pt")
parser.add_argument("--out", default="sessions/flujo_pampar.html")
args = parser.parse_args()
text = args.texto_flag or args.texto
# ── Cargar modelo ─────────────────────────────────────────────────────────
print(f"Cargando modelo desde {args.checkpoint}...")
import sentencepiece as spm
from pampar.coder.v3.config import PRESET_V3
from pampar.coder.v3.modelo import PamparV3
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tok = spm.SentencePieceProcessor()
tok.Load(str(ROOT / "data" / "tokenizer" / "pampar_48k.model"))
model = PamparV3(PRESET_V3).to(device)
ckpt_path = ROOT / args.checkpoint
if ckpt_path.exists():
ckpt = torch.load(str(ckpt_path), map_location=device, weights_only=False)
state = ckpt.get("modelo", ckpt.get("model", ckpt))
model.load_state_dict(state, strict=False)
print(f" Checkpoint cargado: {ckpt_path.name}")
else:
print(f" Checkpoint no encontrado, usando pesos iniciales")
model.registrar_tokenizer(tok)
model.eval()
# ── Forward pass con hooks ────────────────────────────────────────────────
cap = ActivationCapture()
cap.attach(model)
ids = tok.Encode(text)
tokens_str = [tok.IdToPiece(i).replace("▁", " ").strip() or "" for i in ids]
input_ids = torch.tensor([ids], dtype=torch.long, device=device)
# Medir memoria antes
if device.type == "cuda":
torch.cuda.reset_peak_memory_stats()
t0 = time.perf_counter()
with torch.no_grad():
logits, _, _ = model(input_ids)
elapsed_ms = (time.perf_counter() - t0) * 1000
if device.type == "cuda":
mem_mb = torch.cuda.max_memory_allocated() / 1024**2
else:
import os
import psutil
proc = psutil.Process(os.getpid())
mem_mb = proc.memory_info().rss / 1024**2
cap.detach()
# ── Top-5 predicciones del último token ──────────────────────────────────
last_logits = logits[0, -1]
top5 = last_logits.topk(5)
print(f"\nTexto: '{text}'")
print(f"Tokens ({len(ids)}): {tokens_str}")
print(f"Forward pass: {elapsed_ms:.1f}ms | Memoria: {mem_mb:.0f}MB")
print(f"\nTop-5 siguiente token:")
for score, idx in zip(top5.values.tolist(), top5.indices.tolist()):
piece = tok.IdToPiece(idx).replace("▁", " ")
prob = torch.softmax(last_logits, dim=0)[idx].item()
print(f" '{piece}' — {prob * 100:.1f}%")
# ── Generar HTML ──────────────────────────────────────────────────────────
out_path = ROOT / args.out
out_path.parent.mkdir(parents=True, exist_ok=True)
html = build_html(tokens_str, cap, text, elapsed_ms, mem_mb)
out_path.write_text(html, encoding="utf-8")
print(f"\nHTML generado: {out_path}")
print("Abre ese archivo en Chrome/Edge para ver la visualización.")
if __name__ == "__main__":
main()