From 6cb0b47dac1a7f5f7b4e5df2d9b055ea69859c00 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 16 Jun 2026 13:03:07 -0700 Subject: [PATCH 1/4] fix(viewer): restart playback from earliest snapshot --- README.md | 46 ++ analysis/min_angle_graph_plotly_viewer.html | 451 ++++++++++++++++++++ demos/min_angle_graph_export_demo.sh | 68 +++ explorations/min_angle_graph_export.yaml | 39 ++ train.py | 44 ++ train_args.py | 14 + utils/min_angle_graph_export.py | 152 +++++++ 7 files changed, 814 insertions(+) create mode 100644 analysis/min_angle_graph_plotly_viewer.html create mode 100755 demos/min_angle_graph_export_demo.sh create mode 100644 explorations/min_angle_graph_export.yaml create mode 100644 utils/min_angle_graph_export.py diff --git a/README.md b/README.md index 6a8182a68b..fdf66a08a9 100644 --- a/README.md +++ b/README.md @@ -258,6 +258,52 @@ python3 view_model_stats.py run1_stats.csv run2_stats.csv See [documentation/Model_Stats_Table.md](documentation/Model_Stats_Table.md) for more details. +## LM-head minimum-angle graph exports + +`train.py` can optionally export the same closest-neighbor graph used by the +LM-head angle explorer's minimum-angular-distance view after each validation +loss. The export treats each LM-head row as a token vector, streams +row/column blocks through the selected compute device, excludes each token's +self-distance, and records the closest non-self token by signed `0°–180°` +angular distance. The full `vocab_size × vocab_size` angle matrix is never +materialized. + +Enable the export by choosing a labelled output directory and turning on the +per-eval flag: + +```bash +python3 train.py \ + --export_min_angle_graph_dir out/min_angle_graphs \ + --export_min_angle_graph_each_eval \ + --export_min_angle_graph_label shakespeare_char_baseline +``` + +Each validation step writes a labelled CSV and JSON sidecar whose filename +includes the label, iteration, and validation loss: + +```text +out/min_angle_graphs/shakespeare_char_baseline_iter_00000250_val_1.234567.csv +out/min_angle_graphs/shakespeare_char_baseline_iter_00000250_val_1.234567.json +``` + +The CSV is an edge list with one directed nearest-neighbor edge per token: + +```text +token_id,nearest_token_id,min_angle_deg,cosine,token_vector_length,nearest_token_vector_length,min_angle_rank +``` + +Use `--export_min_angle_graph_block_size` to tune temporary matrix multiply +size, and `--export_min_angle_graph_device` to choose `auto`, `cpu`, or a CUDA +device such as `cuda:0`. In `auto` mode, the export uses the LM-head tensor's +current CUDA device when possible, otherwise streams blocks to `cuda:0` if CUDA +is available, and falls back to CPU. + +To review a sequence of exports visually, open +`analysis/min_angle_graph_plotly_viewer.html` in a browser and select the CSV +files from one export directory. The page sorts snapshots by the iteration in +their filenames and provides previous/next/play controls plus a slider to step +from the first validation export to the last. + ## TODO Section: TODO: Add links and descriptions to other Readme's and Demos. diff --git a/analysis/min_angle_graph_plotly_viewer.html b/analysis/min_angle_graph_plotly_viewer.html new file mode 100644 index 0000000000..1b57eda472 --- /dev/null +++ b/analysis/min_angle_graph_plotly_viewer.html @@ -0,0 +1,451 @@ + + + + + + LM-head Minimum-Angle Graph Viewer + + + + +
+

LM-head Minimum-Angle Graph Viewer

+

+ Load one or more CSV exports from --export_min_angle_graph_dir. + The viewer sorts snapshots by iteration and lets you step from the first + validation export to the last without uploading data anywhere. +

+
+ +
+
+
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+
+ + + + +
+
Choose CSV files to begin.
+
+ +
+
+
+
+
+ +
+

Snapshot rows

+
+ + + + + + + + + + + + + +
ranktoken_idnearest_token_idmin_angle_degcosinetoken_vector_lengthnearest_token_vector_length
+
+
+
+ + + + diff --git a/demos/min_angle_graph_export_demo.sh b/demos/min_angle_graph_export_demo.sh new file mode 100755 index 0000000000..0b8fe11037 --- /dev/null +++ b/demos/min_angle_graph_export_demo.sh @@ -0,0 +1,68 @@ +#!/bin/bash + +# Demo the per-validation LM-head minimum-angle graph export on the small +# exploration config, then point the user at the local Plotly viewer. +# +# Optional environment variables: +# OUTPUT_DIR Base output directory for train.py runs. Default: out +# PREFIX Prefix for exploration run names. Default: timestamped demo prefix +# EXPORT_ROOT Directory where the exploration config writes CSV/JSON exports. +# Default: out/min_angle_graph_exports + +OUTPUT_DIR="${OUTPUT_DIR:-out}" +PREFIX="${PREFIX:-min_angle_graph_demo_$(date +%Y%m%d_%H%M%S)_}" +EXPORT_ROOT="${EXPORT_ROOT:-out/min_angle_graph_exports}" +CONFIG="explorations/min_angle_graph_export.yaml" +VIEWER="analysis/min_angle_graph_plotly_viewer.html" + +before_csv_count="$(find "${EXPORT_ROOT}" -type f -name "${PREFIX}*.csv" 2>/dev/null | wc -l | tr -d ' ')" + +cat </dev/null | wc -l | tr -d ' ')" +if [[ "${after_csv_count}" -le "${before_csv_count}" ]]; then + cat >&2 </ + +New CSV exports detected: +$(find "${EXPORT_ROOT}" -type f -name "${PREFIX}*.csv" 2>/dev/null | sort) + +Open the Plotly viewer in your browser: + ${VIEWER} + +Then select one or more exported CSV files from an export directory to step +through validation snapshots from first iteration to last. +EOF diff --git a/explorations/min_angle_graph_export.yaml b/explorations/min_angle_graph_export.yaml new file mode 100644 index 0000000000..dd8bad9818 --- /dev/null +++ b/explorations/min_angle_graph_export.yaml @@ -0,0 +1,39 @@ +# Smoke-test the optional LM-head minimum-angle graph export on a tiny +# validation cadence for both a character-tokenized and GPT-style dataset. +# +# Run with: +# python3 optimization_and_search/run_experiments.py \ +# --config explorations/min_angle_graph_export.yaml \ +# --config_format yaml \ +# --prefix min_angle_export_ +# +# The minipile run has a much larger vocabulary than shakespeare_char, so the +# default smoke settings use CUDA/float16 and a GPU-friendly export block size. +--- + +common_group: + n_layer: [3] + n_head: [3] + n_embd: [150] + block_size: [64] + batch_size: [4] + max_iters: [2000] + eval_interval: [20] + eval_iters: [20] + device: ["cuda:0"] + dtype: ["float16"] + compile: [true] + never_save_checkpoint: [true] + export_min_angle_graph_each_eval: [true] + export_min_angle_graph_device: ["cuda:0"] + +parameter_groups: + - dataset: ["shakespeare_char"] + export_min_angle_graph_block_size: [512] + export_min_angle_graph_label: ["${RUN_NAME}"] + export_min_angle_graph_dir: ["out/min_angle_graph_exports/${RUN_NAME}"] + + - dataset: ["minipile"] + export_min_angle_graph_block_size: [2048] + export_min_angle_graph_label: ["${RUN_NAME}"] + export_min_angle_graph_dir: ["out/min_angle_graph_exports/${RUN_NAME}"] diff --git a/train.py b/train.py index cccf0acba1..4b77dee589 100644 --- a/train.py +++ b/train.py @@ -12,6 +12,22 @@ from collections import deque from datetime import datetime, timedelta +# Some user worktrees may contain helper files named like Python stdlib modules +# (for example copy.py). Preload stdlib modules that Rich/dataclasses need +# before third-party imports so local helper files cannot shadow them during +# interpreter startup. +_repo_dir = os.path.dirname(os.path.abspath(__file__)) +_removed_sys_path_entries = [] +for _entry in ("", _repo_dir): + while _entry in sys.path: + sys.path.remove(_entry) + _removed_sys_path_entries.append(_entry) +import copy as _stdlib_copy +import dataclasses as _stdlib_dataclasses +for _entry in reversed(_removed_sys_path_entries): + sys.path.insert(0, _entry) +del _entry, _removed_sys_path_entries, _repo_dir, _stdlib_copy, _stdlib_dataclasses + from rich.console import Group from rich.console import Console from rich.text import Text @@ -27,6 +43,7 @@ from train_variations.distillation_loss_variants import build_distillation_loss from utils.gpu_monitoring import get_gpu_memory_info, get_process_gpu_memory_bytes +from utils.min_angle_graph_export import export_min_angle_graph as write_min_angle_graph_export from torch.cuda import reset_peak_memory_stats, max_memory_allocated, max_memory_reserved try: @@ -1797,6 +1814,26 @@ def save_checkpoint(self, filename): } torch.save(checkpoint, os.path.join(self.args.out_dir, filename)) + def export_min_angle_graph(self, losses): + """Export the current LM-head minimum-angle graph using the configured writer.""" + export_dir = getattr(self.args, "export_min_angle_graph_dir", None) + if not export_dir: + return + + weight = self.raw_model.apply_lm_head_norm(self.raw_model.lm_head.weight).detach() + label = getattr(self.args, "export_min_angle_graph_label", None) or "min_angle_graph" + val_loss = losses["val"].item() if hasattr(losses["val"], "item") else float(losses["val"]) + csv_path, _ = write_min_angle_graph_export( + weight=weight, + export_dir=export_dir, + label=label, + iter_num=self.iter_num, + val_loss=val_loss, + block_size=getattr(self.args, "export_min_angle_graph_block_size", 2048), + compute_device=getattr(self.args, "export_min_angle_graph_device", "auto"), + ) + print(f"Minimum-angle graph exported to {csv_path}") + def run_validation_step(self, running_mfu, current_epoch, current_dataset, num_steps_with_worse_loss, live=None): losses = self.estimate_loss() @@ -1888,6 +1925,13 @@ def run_validation_step(self, running_mfu, current_epoch, current_dataset, num_s print("Exiting with nan") file.write(str(self.iter_num)) + if self.args.export_min_angle_graph_each_eval: + if live: + live.stop() + self.export_min_angle_graph(losses) + if live: + live.start() + if (not self.args.never_save_checkpoint and self.args.save_major_ckpt_interval is not None): if self.iter_num % self.args.save_major_ckpt_interval == 0: diff --git a/train_args.py b/train_args.py index 835b24db06..424a1c7b39 100644 --- a/train_args.py +++ b/train_args.py @@ -80,6 +80,20 @@ def parse_args(): training_group.add_argument('--log_interval', default=10, type=int) training_group.add_argument('--eval_iters', default=200, type=int) training_group.add_argument('--eval_only', default=False, action=argparse.BooleanOptionalAction) + + # Validation-time LM-head minimum-angle graph export args + training_group.add_argument('--export_min_angle_graph_dir', default=None, type=str, + help='Optional directory for per-validation LM-head minimum-angle graph exports.') + training_group.add_argument('--export_min_angle_graph_each_eval', default=False, + action=argparse.BooleanOptionalAction, + help='Export the LM-head minimum-angle graph after every validation loss.') + training_group.add_argument('--export_min_angle_graph_block_size', default=2048, type=int, + help='Row/column block size used when exporting the minimum-angle graph.') + training_group.add_argument('--export_min_angle_graph_device', default='auto', type=str, + help="Compute device for minimum-angle graph export: 'auto', 'cpu', or a CUDA device such as 'cuda:0'.") + training_group.add_argument('--export_min_angle_graph_label', default=None, type=str, + help='Optional label inserted into minimum-angle graph export filenames.') + training_group.add_argument( '--mezo_epsilon', type=float, diff --git a/utils/min_angle_graph_export.py b/utils/min_angle_graph_export.py new file mode 100644 index 0000000000..d662ad5560 --- /dev/null +++ b/utils/min_angle_graph_export.py @@ -0,0 +1,152 @@ +# utils/min_angle_graph_export.py +import csv +import json +import os + +import torch +import torch.nn.functional as F + + +def resolve_min_angle_graph_device(weight, requested_device="auto"): + """Choose the compute device for blockwise minimum-angle graph export.""" + if requested_device == "auto": + if weight.is_cuda: + return weight.device + if torch.cuda.is_available(): + return torch.device("cuda:0") + return torch.device("cpu") + if requested_device.startswith("cuda") and not torch.cuda.is_available(): + print("CUDA requested for minimum-angle graph export but is unavailable; using CPU.") + return torch.device("cpu") + return torch.device(requested_device) + + +def compute_min_angle_graph(weight, block_size=2048, compute_device="auto"): + """Compute each row vector's closest non-self row by signed angular distance.""" + compute_device = resolve_min_angle_graph_device(weight, compute_device) + block_size = max(1, int(block_size)) + vocab_size = weight.shape[0] + + with torch.no_grad(): + weight = weight.detach() + norms = torch.linalg.vector_norm(weight, ord=2, dim=1).detach().cpu() + best_cosine = torch.full((vocab_size,), -float("inf"), device=compute_device) + best_other_token_id = torch.full((vocab_size,), -1, dtype=torch.long, device=compute_device) + + for row_start in range(0, vocab_size, block_size): + row_end = min(row_start + block_size, vocab_size) + row_block = weight[row_start:row_end].to(compute_device, non_blocking=True) + row_block = F.normalize(row_block, p=2, dim=1) + + for col_start in range(0, vocab_size, block_size): + col_end = min(col_start + block_size, vocab_size) + col_block = weight[col_start:col_end].to(compute_device, non_blocking=True) + col_block = F.normalize(col_block, p=2, dim=1) + cosine_block = row_block @ col_block.T + + if row_start < col_end and col_start < row_end: + diag_start = max(row_start, col_start) + diag_end = min(row_end, col_end) + diag_rows = torch.arange( + diag_start - row_start, + diag_end - row_start, + device=compute_device, + ) + diag_cols = torch.arange( + diag_start - col_start, + diag_end - col_start, + device=compute_device, + ) + cosine_block[diag_rows, diag_cols] = -float("inf") + + block_best_cosine, block_best_col = cosine_block.max(dim=1) + current_best = best_cosine[row_start:row_end] + current_other = best_other_token_id[row_start:row_end] + update_mask = block_best_cosine > current_best + current_best[update_mask] = block_best_cosine[update_mask] + current_other[update_mask] = block_best_col[update_mask] + col_start + + best_cosine = best_cosine.clamp(-1.0, 1.0).cpu() + best_other_token_id = best_other_token_id.cpu() + min_angles = torch.rad2deg(torch.acos(best_cosine)) + sorted_rank = torch.argsort(min_angles) + rank_by_token = torch.empty_like(sorted_rank) + rank_by_token[sorted_rank] = torch.arange(vocab_size, dtype=torch.long) + + return { + "vocab_size": vocab_size, + "block_size": block_size, + "compute_device": str(compute_device), + "norms": norms, + "best_cosine": best_cosine, + "best_other_token_id": best_other_token_id, + "min_angles": min_angles, + "rank_by_token": rank_by_token, + } + + +def safe_filename_label(label): + """Return a filesystem-friendly label while preserving readable separators.""" + return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in label) + + +def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): + """Write a minimum-angle graph CSV plus JSON sidecar and return both paths.""" + export_dir = os.path.expanduser(export_dir) + os.makedirs(export_dir, exist_ok=True) + + safe_label = safe_filename_label(label) + file_stem = f"{safe_label}_iter_{iter_num:08d}_val_{val_loss:.6f}" + csv_path = os.path.join(export_dir, f"{file_stem}.csv") + json_path = os.path.join(export_dir, f"{file_stem}.json") + + norms = graph["norms"] + best_cosine = graph["best_cosine"] + best_other_token_id = graph["best_other_token_id"] + min_angles = graph["min_angles"] + rank_by_token = graph["rank_by_token"] + vocab_size = graph["vocab_size"] + + with open(csv_path, "w", newline="") as csv_file: + writer = csv.writer(csv_file) + writer.writerow([ + "token_id", + "nearest_token_id", + "min_angle_deg", + "cosine", + "token_vector_length", + "nearest_token_vector_length", + "min_angle_rank", + ]) + for token_id in range(vocab_size): + other_id = int(best_other_token_id[token_id]) + writer.writerow([ + token_id, + other_id, + f"{float(min_angles[token_id]):.9f}", + f"{float(best_cosine[token_id]):.9f}", + f"{float(norms[token_id]):.9f}", + f"{float(norms[other_id]):.9f}" if other_id >= 0 else "", + int(rank_by_token[token_id]), + ]) + + metadata = { + "iter_num": iter_num, + "val_loss": val_loss, + "label": label, + "csv_path": csv_path, + "vocab_size": vocab_size, + "block_size": graph["block_size"], + "compute_device": graph["compute_device"], + "angle_definition": "signed 0-180 degrees, closest non-self token by maximum cosine", + } + with open(json_path, "w") as json_file: + json.dump(metadata, json_file, indent=2) + + return csv_path, json_path + + +def export_min_angle_graph(weight, export_dir, label, iter_num, val_loss, block_size=2048, compute_device="auto"): + """Compute and write the LM-head minimum-angle graph export.""" + graph = compute_min_angle_graph(weight, block_size=block_size, compute_device=compute_device) + return write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss) From 7d54e410cf166ace4f604f6b20198223a437522c Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 16 Jun 2026 23:11:51 -0700 Subject: [PATCH 2/4] feat(viewer): show escaped token strings in min angle exports --- README.md | 2 +- analysis/min_angle_graph_plotly_viewer.html | 68 ++++++++++++++++++--- train.py | 5 ++ utils/min_angle_graph_export.py | 39 +++++++++++- 4 files changed, 102 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index fdf66a08a9..3049cd48ac 100644 --- a/README.md +++ b/README.md @@ -289,7 +289,7 @@ out/min_angle_graphs/shakespeare_char_baseline_iter_00000250_val_1.234567.json The CSV is an edge list with one directed nearest-neighbor edge per token: ```text -token_id,nearest_token_id,min_angle_deg,cosine,token_vector_length,nearest_token_vector_length,min_angle_rank +token_id,token_text_escaped,nearest_token_id,nearest_token_text_escaped,min_angle_deg,cosine,token_vector_length,nearest_token_vector_length,min_angle_rank ``` Use `--export_min_angle_graph_block_size` to tune temporary matrix multiply diff --git a/analysis/min_angle_graph_plotly_viewer.html b/analysis/min_angle_graph_plotly_viewer.html index 1b57eda472..1fd093f205 100644 --- a/analysis/min_angle_graph_plotly_viewer.html +++ b/analysis/min_angle_graph_plotly_viewer.html @@ -191,8 +191,10 @@

Snapshot rows

rank token_id + token_text_escaped nearest_token_id - min_angle_deg + nearest_token_text_escaped + min_angle_deg / magnitudes cosine token_vector_length nearest_token_vector_length @@ -232,14 +234,37 @@

Snapshot rows

}; } + function parseCsvLine(line) { + const values = []; + let value = ""; + let inQuotes = false; + for (let index = 0; index < line.length; index += 1) { + const char = line[index]; + const nextChar = line[index + 1]; + if (char === '"' && inQuotes && nextChar === '"') { + value += '"'; + index += 1; + } else if (char === '"') { + inQuotes = !inQuotes; + } else if (char === "," && !inQuotes) { + values.push(value); + value = ""; + } else { + value += char; + } + } + values.push(value); + return values; + } + function parseCsv(text) { const lines = text.trim().split(/\r?\n/).filter(Boolean); if (lines.length < 2) { return []; } - const headers = lines[0].split(",").map((h) => h.trim()); + const headers = parseCsvLine(lines[0]).map((h) => h.trim()); return lines.slice(1).map((line) => { - const values = line.split(","); + const values = parseCsvLine(line); const row = {}; headers.forEach((header, index) => { const raw = values[index] ?? ""; @@ -250,6 +275,15 @@

Snapshot rows

}); } + function escapedCell(value) { + return String(value ?? "") + .replaceAll("&", "&") + .replaceAll("<", "<") + .replaceAll(">", ">") + .replaceAll('"', """) + .replaceAll("'", "'"); + } + function readFile(file) { return new Promise((resolve, reject) => { const reader = new FileReader(); @@ -312,8 +346,16 @@

Snapshot rows

x: rows.map((row) => row.min_angle_rank), y: rows.map((row) => row.min_angle_deg), marker: {size: 5, color: colors, opacity: 0.75}, - customdata: rows.map((row) => [row.token_id, row.nearest_token_id, row.cosine]), - hovertemplate: "rank=%{x}
angle=%{y:.6f}°
token=%{customdata[0]}
nearest=%{customdata[1]}
cosine=%{customdata[2]:.6f}", + customdata: rows.map((row) => [ + row.token_id, + row.token_text_escaped, + row.nearest_token_id, + row.nearest_token_text_escaped, + row.cosine, + row.token_vector_length, + row.nearest_token_vector_length, + ]), + hovertemplate: "rank=%{x}
angle=%{y:.6f}°
token=%{customdata[0]} (%{customdata[1]})
nearest=%{customdata[2]} (%{customdata[3]})
magnitudes=%{customdata[5]:.6f} → %{customdata[6]:.6f}
cosine=%{customdata[4]:.6f}", }], { ...commonLayout, title: "Minimum angle by sorted rank", @@ -346,8 +388,16 @@

Snapshot rows

colorbar: {title: "angle°"}, opacity: 0.75, }, - customdata: rows.map((row) => [row.min_angle_rank, row.cosine]), - hovertemplate: "token=%{x}
nearest=%{y}
rank=%{customdata[0]}
cosine=%{customdata[1]:.6f}", + customdata: rows.map((row) => [ + row.min_angle_rank, + row.cosine, + row.token_text_escaped, + row.nearest_token_text_escaped, + row.min_angle_deg, + row.token_vector_length, + row.nearest_token_vector_length, + ]), + hovertemplate: "token=%{x} (%{customdata[2]})
nearest=%{y} (%{customdata[3]})
rank=%{customdata[0]}
angle=%{customdata[4]:.6f}°
magnitudes=%{customdata[5]:.6f} → %{customdata[6]:.6f}
cosine=%{customdata[1]:.6f}", }], { ...commonLayout, title: "Nearest-neighbor edge map", @@ -359,8 +409,10 @@

Snapshot rows

${row.min_angle_rank} ${row.token_id} + ${escapedCell(row.token_text_escaped)} ${row.nearest_token_id} - ${Number(row.min_angle_deg).toFixed(6)} + ${escapedCell(row.nearest_token_text_escaped)} + ${Number(row.min_angle_deg).toFixed(6)}° / ${Number(row.token_vector_length).toFixed(6)} → ${Number(row.nearest_token_vector_length).toFixed(6)} ${Number(row.cosine).toFixed(6)} ${Number(row.token_vector_length).toFixed(6)} ${Number(row.nearest_token_vector_length).toFixed(6)} diff --git a/train.py b/train.py index 4b77dee589..7e589e7161 100644 --- a/train.py +++ b/train.py @@ -1821,6 +1821,10 @@ def export_min_angle_graph(self, losses): return weight = self.raw_model.apply_lm_head_norm(self.raw_model.lm_head.weight).detach() + if not hasattr(self, "_min_angle_graph_token_texts"): + self._min_angle_graph_token_texts = [ + self.decode([token_id]) for token_id in range(weight.shape[0]) + ] label = getattr(self.args, "export_min_angle_graph_label", None) or "min_angle_graph" val_loss = losses["val"].item() if hasattr(losses["val"], "item") else float(losses["val"]) csv_path, _ = write_min_angle_graph_export( @@ -1831,6 +1835,7 @@ def export_min_angle_graph(self, losses): val_loss=val_loss, block_size=getattr(self.args, "export_min_angle_graph_block_size", 2048), compute_device=getattr(self.args, "export_min_angle_graph_device", "auto"), + token_texts=self._min_angle_graph_token_texts, ) print(f"Minimum-angle graph exported to {csv_path}") diff --git a/utils/min_angle_graph_export.py b/utils/min_angle_graph_export.py index d662ad5560..fe9d92c1ac 100644 --- a/utils/min_angle_graph_export.py +++ b/utils/min_angle_graph_export.py @@ -90,7 +90,12 @@ def safe_filename_label(label): return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in label) -def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): +def escape_token_text(text): + """Escape token text for CSV/HTML display while preserving visible escapes.""" + return str(text).encode("unicode_escape").decode("ascii") + + +def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss, token_texts=None): """Write a minimum-angle graph CSV plus JSON sidecar and return both paths.""" export_dir = os.path.expanduser(export_dir) os.makedirs(export_dir, exist_ok=True) @@ -106,12 +111,21 @@ def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): min_angles = graph["min_angles"] rank_by_token = graph["rank_by_token"] vocab_size = graph["vocab_size"] + escaped_token_texts = None + if token_texts is not None: + if len(token_texts) != vocab_size: + raise ValueError( + f"Expected {vocab_size} token strings, got {len(token_texts)}." + ) + escaped_token_texts = [escape_token_text(text) for text in token_texts] with open(csv_path, "w", newline="") as csv_file: writer = csv.writer(csv_file) writer.writerow([ "token_id", + "token_text_escaped", "nearest_token_id", + "nearest_token_text_escaped", "min_angle_deg", "cosine", "token_vector_length", @@ -122,7 +136,9 @@ def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): other_id = int(best_other_token_id[token_id]) writer.writerow([ token_id, + escaped_token_texts[token_id] if escaped_token_texts is not None else "", other_id, + escaped_token_texts[other_id] if escaped_token_texts is not None and other_id >= 0 else "", f"{float(min_angles[token_id]):.9f}", f"{float(best_cosine[token_id]):.9f}", f"{float(norms[token_id]):.9f}", @@ -138,6 +154,7 @@ def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): "vocab_size": vocab_size, "block_size": graph["block_size"], "compute_device": graph["compute_device"], + "has_token_text_escaped": escaped_token_texts is not None, "angle_definition": "signed 0-180 degrees, closest non-self token by maximum cosine", } with open(json_path, "w") as json_file: @@ -146,7 +163,23 @@ def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): return csv_path, json_path -def export_min_angle_graph(weight, export_dir, label, iter_num, val_loss, block_size=2048, compute_device="auto"): +def export_min_angle_graph( + weight, + export_dir, + label, + iter_num, + val_loss, + block_size=2048, + compute_device="auto", + token_texts=None, +): """Compute and write the LM-head minimum-angle graph export.""" graph = compute_min_angle_graph(weight, block_size=block_size, compute_device=compute_device) - return write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss) + return write_min_angle_graph_export( + graph, + export_dir, + label, + iter_num, + val_loss, + token_texts=token_texts, + ) From bdfb6e88632207e1f26c2334c7ccf4cb93dbb5e8 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 16 Jun 2026 23:27:29 -0700 Subject: [PATCH 3/4] fix(viewer): render token strings in hover text --- analysis/min_angle_graph_plotly_viewer.html | 42 +++++++++++---------- 1 file changed, 22 insertions(+), 20 deletions(-) diff --git a/analysis/min_angle_graph_plotly_viewer.html b/analysis/min_angle_graph_plotly_viewer.html index 1fd093f205..764bf68f61 100644 --- a/analysis/min_angle_graph_plotly_viewer.html +++ b/analysis/min_angle_graph_plotly_viewer.html @@ -325,6 +325,24 @@

Snapshot rows

return rows.slice(0, Number(maxRowsInput.value || 250)); } + function formatNumber(value, digits = 6) { + const numberValue = Number(value); + return Number.isFinite(numberValue) ? numberValue.toFixed(digits) : ""; + } + + function formatRowHover(row, rankLabel) { + const tokenText = escapedCell(row.token_text_escaped || ""); + const nearestText = escapedCell(row.nearest_token_text_escaped || ""); + return [ + `${rankLabel}=${escapedCell(row.min_angle_rank)}`, + `angle=${formatNumber(row.min_angle_deg)}°`, + `token=${escapedCell(row.token_id)} (${tokenText})`, + `nearest=${escapedCell(row.nearest_token_id)} (${nearestText})`, + `magnitudes=${formatNumber(row.token_vector_length)} → ${formatNumber(row.nearest_token_vector_length)}`, + `cosine=${formatNumber(row.cosine)}`, + ].join("
"); + } + function plotSnapshot() { if (!state.snapshots.length) { return; @@ -346,16 +364,8 @@

Snapshot rows

x: rows.map((row) => row.min_angle_rank), y: rows.map((row) => row.min_angle_deg), marker: {size: 5, color: colors, opacity: 0.75}, - customdata: rows.map((row) => [ - row.token_id, - row.token_text_escaped, - row.nearest_token_id, - row.nearest_token_text_escaped, - row.cosine, - row.token_vector_length, - row.nearest_token_vector_length, - ]), - hovertemplate: "rank=%{x}
angle=%{y:.6f}°
token=%{customdata[0]} (%{customdata[1]})
nearest=%{customdata[2]} (%{customdata[3]})
magnitudes=%{customdata[5]:.6f} → %{customdata[6]:.6f}
cosine=%{customdata[4]:.6f}", + text: rows.map((row) => formatRowHover(row, "rank")), + hovertemplate: "%{text}", }], { ...commonLayout, title: "Minimum angle by sorted rank", @@ -388,16 +398,8 @@

Snapshot rows

colorbar: {title: "angle°"}, opacity: 0.75, }, - customdata: rows.map((row) => [ - row.min_angle_rank, - row.cosine, - row.token_text_escaped, - row.nearest_token_text_escaped, - row.min_angle_deg, - row.token_vector_length, - row.nearest_token_vector_length, - ]), - hovertemplate: "token=%{x} (%{customdata[2]})
nearest=%{y} (%{customdata[3]})
rank=%{customdata[0]}
angle=%{customdata[4]:.6f}°
magnitudes=%{customdata[5]:.6f} → %{customdata[6]:.6f}
cosine=%{customdata[1]:.6f}", + text: rows.map((row) => formatRowHover(row, "rank")), + hovertemplate: "%{text}", }], { ...commonLayout, title: "Nearest-neighbor edge map", From 23b2241fb46aab89cdc10cd6a9b507628e183a98 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 16 Jun 2026 23:36:24 -0700 Subject: [PATCH 4/4] fix(viewer): warn when token text is missing --- analysis/min_angle_graph_plotly_viewer.html | 22 ++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/analysis/min_angle_graph_plotly_viewer.html b/analysis/min_angle_graph_plotly_viewer.html index 764bf68f61..e2cc0d3285 100644 --- a/analysis/min_angle_graph_plotly_viewer.html +++ b/analysis/min_angle_graph_plotly_viewer.html @@ -284,6 +284,14 @@

Snapshot rows

.replaceAll("'", "'"); } + function tokenTextOrMissing(row, textKey) { + const value = row[textKey]; + if (value === undefined || value === null || value === "") { + return "[token text missing: regenerate CSV exports]"; + } + return value; + } + function readFile(file) { return new Promise((resolve, reject) => { const reader = new FileReader(); @@ -309,7 +317,11 @@

Snapshot rows

function snapshotLabel(snapshot, index) { const iter = Number.isFinite(snapshot.iter) ? snapshot.iter : "unknown"; const valLoss = Number.isFinite(snapshot.valLoss) ? snapshot.valLoss.toFixed(6) : "unknown"; - return `Snapshot ${index + 1}/${state.snapshots.length}: ${snapshot.name}\niteration=${iter}, val_loss=${valLoss}, rows=${snapshot.rows.length}`; + const hasTokenText = snapshot.rows.some((row) => row.token_text_escaped || row.nearest_token_text_escaped); + const tokenTextNote = hasTokenText + ? "" + : "\nToken text columns are empty/missing; regenerate CSV exports with the latest exporter to show actual token strings."; + return `Snapshot ${index + 1}/${state.snapshots.length}: ${snapshot.name}\niteration=${iter}, val_loss=${valLoss}, rows=${snapshot.rows.length}${tokenTextNote}`; } function currentRowsForTable(snapshot) { @@ -331,8 +343,8 @@

Snapshot rows

} function formatRowHover(row, rankLabel) { - const tokenText = escapedCell(row.token_text_escaped || ""); - const nearestText = escapedCell(row.nearest_token_text_escaped || ""); + const tokenText = escapedCell(tokenTextOrMissing(row, "token_text_escaped")); + const nearestText = escapedCell(tokenTextOrMissing(row, "nearest_token_text_escaped")); return [ `${rankLabel}=${escapedCell(row.min_angle_rank)}`, `angle=${formatNumber(row.min_angle_deg)}°`, @@ -411,9 +423,9 @@

Snapshot rows

${row.min_angle_rank} ${row.token_id} - ${escapedCell(row.token_text_escaped)} + ${escapedCell(tokenTextOrMissing(row, "token_text_escaped"))} ${row.nearest_token_id} - ${escapedCell(row.nearest_token_text_escaped)} + ${escapedCell(tokenTextOrMissing(row, "nearest_token_text_escaped"))} ${Number(row.min_angle_deg).toFixed(6)}° / ${Number(row.token_vector_length).toFixed(6)} → ${Number(row.nearest_token_vector_length).toFixed(6)} ${Number(row.cosine).toFixed(6)} ${Number(row.token_vector_length).toFixed(6)}