""" Visualization and analysis (Person E). Tasks: training curves (TensorBoard or training_summary.json), experiment comparison plots, cross-attention heatmaps, translation-example Markdown. Examples: python scripts/visualize.py --task training_curves --log-dir outputs/exp1/logs python scripts/visualize.py --task training_curves --summary-json checkpoints/training_summary.json python scripts/visualize.py --task comparison --results-dir outputs/experiments python scripts/visualize.py --task attention --config configs/default_config.yaml \\ --checkpoint checkpoints/best_model.pt --src "Hello ." --tgt "Hi there ." python scripts/visualize.py --task examples --config configs/default_config.yaml \\ --checkpoint checkpoints/best_model.pt --pairs-json result/translation_pairs.example.json """ from __future__ import annotations import argparse import json import logging import math import sys from pathlib import Path from typing import Any, Optional import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch import yaml from omegaconf import OmegaConf sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) REPO_ROOT = Path(__file__).resolve().parent.parent logger = logging.getLogger(__name__) plt.rcParams["axes.unicode_minus"] = False def _format_metrics_table(labels: list[str], keys: list[str], metrics_map: dict[str, list[float]]) -> str: """Space-padded table (no tabs; Matplotlib renders tabs poorly).""" header = ["Experiment"] + [k.upper() for k in keys] rows: list[list[str]] = [header] for i, lab in enumerate(labels): rows.append( [str(lab)] + [f"{metrics_map[k][i]:.4f}" if i < len(metrics_map[k]) else "-" for k in keys] ) ncols = len(header) widths = [max(len(rows[r][c]) for r in range(len(rows))) for c in range(ncols)] out_lines = [] for row in rows: out_lines.append(" ".join(row[c].ljust(widths[c]) for c in range(ncols))) return "\n".join(out_lines) def _read_training_summary(path: Path) -> dict[str, Any]: with open(path, "r", encoding="utf-8") as f: return json.load(f) def _read_tensorboard_scalars(log_dir: Path) -> dict[str, tuple[list[int], list[float]]]: try: from tensorboard.backend.event_processing.event_accumulator import EventAccumulator except ImportError as e: raise ImportError( "TensorBoard is required for --log-dir. Install with: pip install tensorboard" ) from e series: dict[str, tuple[list[int], list[float]]] = {} log_dir = Path(log_dir) if not log_dir.exists(): return series ea = EventAccumulator(str(log_dir), size_guidance={"scalars": 0}) ea.Reload() for tag in ea.Tags().get("scalars", []): events = ea.Scalars(tag) steps = [e.step for e in events] vals = [e.value for e in events] series[tag] = (steps, vals) return series def plot_training_curves( log_dir: Optional[str] = None, output_path: str = "outputs/training_curves.png", summary_json: Optional[str] = None, ) -> Path: """Plot training curves from training_summary.json or TensorBoard scalars.""" out = Path(output_path) out.parent.mkdir(parents=True, exist_ok=True) fig, axes = plt.subplots(2, 2, figsize=(12, 8)) if summary_json: summ_path = Path(summary_json) if not summ_path.is_file(): raise FileNotFoundError(f"training_summary.json not found: {summ_path}") data = _read_training_summary(summ_path) epochs = list(range(1, len(data.get("train_loss_history", [])) + 1)) tl = data.get("train_loss_history", []) axes[0, 0].plot(epochs, tl, marker="o") axes[0, 0].set_title("Train Loss (per epoch)") axes[0, 0].set_xlabel("Epoch") axes[0, 0].set_ylabel("Loss") axes[0, 0].grid(True, alpha=0.3) vm = data.get("val_metrics_history", []) if vm: val_loss = [m.get("val_loss", float("nan")) for m in vm] axes[0, 1].plot(range(1, len(val_loss) + 1), val_loss, marker="o", color="tab:orange") axes[0, 1].set_title("Val Loss") axes[0, 1].set_xlabel("Epoch") axes[0, 1].grid(True, alpha=0.3) bleu = [m.get("bleu") for m in vm if isinstance(m.get("bleu"), (int, float))] if bleu: axes[1, 0].plot(range(1, len(bleu) + 1), bleu, marker="o", color="tab:green") axes[1, 0].set_title("BLEU (validation)") axes[1, 0].set_xlabel("Epoch") axes[1, 0].grid(True, alpha=0.3) else: axes[1, 0].text(0.5, 0.5, "No BLEU in validation logs", ha="center", va="center") axes[1, 0].axis("off") axes[1, 1].text( 0.1, 0.5, f"best_epoch: {data.get('best_epoch')}\n" f"metric: {data.get('metric_name')}\n" f"best: {data.get('best_metric')}\n" f"steps: {data.get('total_steps')}", fontsize=11, va="center", ) axes[1, 1].axis("off") axes[1, 1].set_title("Summary") elif log_dir: series = _read_tensorboard_scalars(Path(log_dir)) if not series: raise RuntimeError(f"No TensorBoard scalar events under {log_dir!r}") def plot_tag(ax, tag: str, title: str): if tag not in series: return steps, vals = series[tag] ax.plot(steps, vals) ax.set_title(title) ax.set_xlabel("Step") ax.grid(True, alpha=0.3) plot_tag(axes[0, 0], "Loss/train_step", "Train Loss (step)") plot_tag(axes[0, 1], "Loss/train", "Train Loss (epoch)") plot_tag(axes[1, 0], "Metrics/bleu", "BLEU") plot_tag(axes[1, 1], "LR/step", "Learning Rate") else: raise ValueError("Provide either --summary-json or --log-dir") fig.suptitle("EasyTranslate Training Curves", fontsize=14) fig.tight_layout() fig.savefig(out, dpi=150) plt.close(fig) logger.info("Saved training curves: %s", out) return out def _collect_experiment_metrics(results_dir: Path) -> tuple[list[str], dict[str, list[float]]]: """Load metrics from experiments_summary.json or per-run evaluation_results.json under subdirs.""" results_dir = Path(results_dir) labels: list[str] = [] metrics_map: dict[str, list[float]] = {} direct = results_dir / "evaluation_results.json" if direct.is_file(): labels.append(results_dir.name or "single") with open(direct, "r", encoding="utf-8") as f: m = json.load(f) for k, v in m.items(): if isinstance(v, (int, float)) and not isinstance(v, bool): metrics_map.setdefault(k, []).append(float(v)) return labels, metrics_map summary_file = results_dir / "experiments_summary.json" if summary_file.is_file(): with open(summary_file, "r", encoding="utf-8") as f: rows = json.load(f) for row in rows: name = row.get("name", "unknown") labels.append(name) m = row.get("metrics") or {} for k, v in m.items(): if isinstance(v, (int, float)) and not isinstance(v, bool): metrics_map.setdefault(k, []).append(float(v)) return labels, metrics_map for sub in sorted(results_dir.iterdir()): if not sub.is_dir(): continue ev = sub / "evaluation_results.json" if not ev.is_file(): continue labels.append(sub.name) with open(ev, "r", encoding="utf-8") as f: m = json.load(f) for k, v in m.items(): if isinstance(v, (int, float)) and not isinstance(v, bool): metrics_map.setdefault(k, []).append(float(v)) if len(labels) != len(next(iter(metrics_map.values()), [])) and metrics_map: # Metric length mismatch across runs: keep rows; plotting filters by available keys. pass return labels, metrics_map def plot_experiment_comparison( results_dir: str, output_path: str = "outputs/experiment_comparison.png", ) -> Path: """Bar chart for BLEU / COMET / chrF / etc., plus a small text table.""" out = Path(output_path) out.parent.mkdir(parents=True, exist_ok=True) labels, metrics_map = _collect_experiment_metrics(Path(results_dir)) if not labels: raise RuntimeError( f"No experiments_summary.json or */evaluation_results.json under {results_dir}" ) preferred = ["bleu", "comet", "chrf", "ter"] keys = [k for k in preferred if k in metrics_map and len(metrics_map[k]) == len(labels)] if not keys: keys = [k for k, vals in metrics_map.items() if len(vals) == len(labels)] if not keys: raise RuntimeError( "No numeric metric columns aligned with each experiment; check evaluation_results.json" ) n = len(keys) fig, axes = plt.subplots(1, max(n, 1), figsize=(4 * max(n, 1), 4)) if n == 1: axes = [axes] for ax, key in zip(axes, keys): vals = metrics_map[key][: len(labels)] x = np.arange(len(labels)) ax.bar(x, vals, color="steelblue") ax.set_xticks(x) ax.set_xticklabels(labels, rotation=25, ha="right") ax.set_title(key.upper()) ax.grid(True, axis="y", alpha=0.3) table_text = _format_metrics_table(labels, keys, metrics_map) fig.subplots_adjust(bottom=0.28) fig.text(0.04, 0.02, table_text, fontsize=9, va="bottom", ha="left") fig.suptitle("Experiment Comparison", fontsize=14) fig.tight_layout() fig.savefig(out, dpi=150, bbox_inches="tight", pad_inches=0.25) plt.close(fig) logger.info("Saved experiment comparison plot: %s", out) return out def _cross_attention_weight_matrix( attn_module: torch.nn.Module, query: torch.Tensor, key: torch.Tensor, memory_key_padding_mask: Optional[torch.BoolTensor], ) -> torch.Tensor: """Scaled dot-product attention weights [B, L_q, L_k], head-mean (for Flash / standard MHAttention).""" B, L_q, _ = query.shape L_k = key.shape[1] nhead = attn_module.nhead d_k = attn_module.d_k Q = attn_module.q_proj(query).view(B, L_q, nhead, d_k).transpose(1, 2) K = attn_module.k_proj(key).view(B, L_k, nhead, d_k).transpose(1, 2) if getattr(attn_module, "rope", None) is not None and attn_module.rope is not None: Q, K = attn_module.rope.apply_rotary_pos_emb(Q, K) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if memory_key_padding_mask is not None: scores = scores.masked_fill( memory_key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf"), ) w = torch.softmax(scores, dim=-1).mean(dim=1) return w[0] def visualize_attention( model: torch.nn.Module, src_text: str, tgt_text: str, tokenizer, output_path: str = "outputs/attention_map.png", layer_idx: int = -1, ) -> Path: """ Cross-attention alignment heatmap for the last (or chosen) decoder layer. Only models with ``decoder.layers[*].multihead_attn`` (e.g. TransformerTranslationModel). """ from easytranslate.model.transformer import TransformerTranslationModel if not isinstance(model, TransformerTranslationModel): raise TypeError("visualize_attention only supports TransformerTranslationModel") device = next(model.parameters()).device model.eval() src_ids_list = tokenizer.encode(src_text, add_special_tokens=True) tgt_ids_list = tokenizer.encode(tgt_text, add_special_tokens=True) if len(tgt_ids_list) < 2: raise ValueError("target sequence too short for teacher-forcing visualization") teacher_tgt = tgt_ids_list[:-1] src_ids = torch.tensor([src_ids_list], dtype=torch.long, device=device) tgt_in = torch.tensor([teacher_tgt], dtype=torch.long, device=device) pad_id = model.pad_id src_padding = src_ids.eq(pad_id) tgt_padding = tgt_in.eq(pad_id) captured: dict[str, Any] = {} layer = model.decoder.layers[layer_idx] def _hook_layer_kw(m, args, kwargs, output): tgt_side, memory = args[0], args[1] mem_pad = kwargs.get("memory_key_padding_mask") query = m.norm2(tgt_side) captured["weights"] = _cross_attention_weight_matrix( m.multihead_attn, query, memory, mem_pad ) def _hook_layer_legacy(m, inp, output): tgt_side, memory = inp[0], inp[1] query = m.norm2(tgt_side) captured["weights"] = _cross_attention_weight_matrix( m.multihead_attn, query, memory, None ) try: handle = layer.register_forward_hook(_hook_layer_kw, with_kwargs=True) except TypeError: handle = layer.register_forward_hook(_hook_layer_legacy) with torch.no_grad(): logits = model(src_ids, tgt_in, src_padding, tgt_padding) handle.remove() if "weights" not in captured: raise RuntimeError("cross-attention hook did not run") w = captured["weights"].detach().float().cpu().numpy() _ = logits src_tokens = [tokenizer.decode([i]) for i in src_ids_list] tgt_tokens = [tokenizer.decode([i]) for i in teacher_tgt] fig, ax = plt.subplots(figsize=(max(8, w.shape[1] * 0.35), max(6, w.shape[0] * 0.35))) im = ax.imshow(w, cmap="viridis", aspect="auto") ax.set_xticks(range(len(src_tokens))) ax.set_yticks(range(len(tgt_tokens))) ax.set_xticklabels(src_tokens, rotation=45, ha="right", fontsize=8) ax.set_yticklabels(tgt_tokens, fontsize=8) ax.set_xlabel("Source") ax.set_ylabel("Target (teacher forcing)") ax.set_title("Cross-attention (last layer, heads mean)") fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) out = Path(output_path) out.parent.mkdir(parents=True, exist_ok=True) fig.tight_layout() fig.savefig(out, dpi=150) plt.close(fig) logger.info("Saved attention heatmap: %s", out) return out def generate_translation_examples( evaluator, test_pairs: list[tuple[str, str]], output_path: str = "outputs/translation_examples.md", ) -> Path: """Write Markdown: source, reference, hypothesis, sentence BLEU and chrF.""" from easytranslate.evaluation.metrics import compute_bleu, compute_chrf srcs = [p[0] for p in test_pairs] refs = [p[1] for p in test_pairs] hyps = evaluator.translate(srcs) def esc(t: str) -> str: return t.replace("|", "\\|").replace("\n", " ") lines = [ "# Translation examples", "", "| # | Source | Reference | Hypothesis | sent-BLEU | sent-chrF |", "|---|--------|-----------|------------|-----------|-----------|", ] for i, (s, r, h) in enumerate(zip(srcs, refs, hyps), 1): sb = compute_bleu([h], [r])["bleu"] ch = compute_chrf([h], [r])["chrf"] lines.append(f"| {i} | {esc(s)} | {esc(r)} | {esc(h)} | {sb:.2f} | {ch:.2f} |") lines.append("") lines.append("> Sentence BLEU/chrF are indicative only (tokenization-dependent).") out = Path(output_path) out.parent.mkdir(parents=True, exist_ok=True) out.write_text("\n".join(lines), encoding="utf-8") logger.info("Wrote translation examples: %s", out) return out def _load_model_for_visual( config_path: Path, checkpoint_path: Path, ) -> tuple[torch.nn.Module, Any, dict]: """Load scratch Transformer + tokenizer from YAML and checkpoint (prefers config inside checkpoint).""" with open(config_path, "r", encoding="utf-8") as f: file_cfg = yaml.safe_load(f) try: ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) except TypeError: ckpt = torch.load(checkpoint_path, map_location="cpu") cfg = ckpt.get("config") if cfg is None: cfg = file_cfg else: try: from omegaconf import DictConfig if isinstance(cfg, DictConfig): cfg = OmegaConf.to_container(cfg, resolve=True) except Exception: pass if not isinstance(cfg, dict): cfg = dict(cfg) from easytranslate.model.transformer import TransformerTranslationModel from easytranslate.data.tokenizer import build_tokenizer tok_cfg = cfg.get("tokenizer") or cfg.get("data", {}).get("tokenizer") or file_cfg.get("tokenizer") or {} try: tokenizer = build_tokenizer(tok_cfg) except ValueError as e: raise ValueError( "Cannot build tokenizer: set tokenizer.path in config or store a loadable tokenizer " "section in checkpoint['config']." ) from e mcfg = cfg.get("model", {}).get("transformer", {}) or file_cfg.get("model", {}).get("transformer", {}) model = TransformerTranslationModel( src_vocab_size=tokenizer.vocab_size, tgt_vocab_size=tokenizer.vocab_size, d_model=int(mcfg.get("d_model", 512)), nhead=int(mcfg.get("nhead", 8)), num_encoder_layers=int(mcfg.get("num_encoder_layers", 6)), num_decoder_layers=int(mcfg.get("num_decoder_layers", 6)), dim_feedforward=int(mcfg.get("dim_feedforward", 2048)), dropout=float(mcfg.get("dropout", 0.1)), activation=str(mcfg.get("activation", "gelu")), max_seq_len=int(mcfg.get("max_seq_len", 512)), use_flash_attention=bool(mcfg.get("use_flash_attention", True)), use_rotary_embedding=bool(mcfg.get("use_rotary_embedding", True)), pre_norm=bool(mcfg.get("pre_norm", True)), pad_id=tokenizer.pad_token_id, ) model.load_state_dict(ckpt["model_state_dict"]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) return model, tokenizer, cfg def _parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="EasyTranslate visualization CLI") p.add_argument( "--task", choices=["training_curves", "comparison", "attention", "examples"], required=True, ) p.add_argument("--log-dir", type=str, default=None) p.add_argument("--summary-json", type=str, default=None) p.add_argument("--results-dir", type=str, default=None) p.add_argument("--output", type=str, default=None) p.add_argument("--config", type=str, default="configs/default_config.yaml") p.add_argument("--checkpoint", type=str, default=None) p.add_argument("--src", type=str, default=None) p.add_argument("--tgt", type=str, default=None) p.add_argument("--pairs-json", type=str, default=None) return p.parse_args() def main() -> None: logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") args = _parse_args() if args.task == "training_curves": outp = args.output or "outputs/training_curves.png" plot_training_curves( log_dir=args.log_dir, output_path=outp, summary_json=args.summary_json, ) print(f"OK: {outp}") elif args.task == "comparison": rd = args.results_dir or "outputs/experiments" outp = args.output or "outputs/experiment_comparison.png" plot_experiment_comparison(rd, outp) print(f"OK: {outp}") elif args.task == "attention": if not args.checkpoint or not args.src or not args.tgt: raise SystemExit("--task attention requires --checkpoint --src --tgt") cfg_p = (REPO_ROOT / args.config).resolve() ckpt_p = (REPO_ROOT / args.checkpoint).resolve() model, tokenizer, _ = _load_model_for_visual(cfg_p, ckpt_p) outp = args.output or "outputs/attention_map.png" visualize_attention(model, args.src, args.tgt, tokenizer, output_path=outp) print(f"OK: {outp}") elif args.task == "examples": if not args.checkpoint or not args.pairs_json: raise SystemExit("--task examples requires --checkpoint --pairs-json") cfg_p = (REPO_ROOT / args.config).resolve() ckpt_p = (REPO_ROOT / args.checkpoint).resolve() model, tokenizer, cfg = _load_model_for_visual(cfg_p, ckpt_p) from easytranslate.evaluation.evaluator import Evaluator evaluator = Evaluator(model=model, tokenizer=tokenizer, config=cfg) pairs_path = (REPO_ROOT / args.pairs_json).resolve() with open(pairs_path, "r", encoding="utf-8") as f: raw = json.load(f) pairs: list[tuple[str, str]] = [] for item in raw: if isinstance(item, dict): pairs.append((item["src"], item["ref"])) else: pairs.append((item[0], item[1])) outp = args.output or "outputs/translation_examples.md" generate_translation_examples(evaluator, pairs, outp) print(f"OK: {outp}") if __name__ == "__main__": main()