import argparse import json from pathlib import Path import matplotlib.pyplot as plt import numpy as np DATA_DIR = Path("data/personas") OUTPUT_PATH = Path("artifacts/embedding_comparison.png") def load_runs(): runs = [] for path in DATA_DIR.glob("*.json"): try: run = json.loads(path.read_text(encoding="utf-8")) except Exception: continue answers = run.get("answers", []) if not answers: continue dim = len(answers[0].get("embedding", [])) if dim: run["_path"] = str(path) run["_dim"] = dim runs.append(run) return sorted(runs, key=lambda item: item.get("created_at", ""), reverse=True) def pick_runs(runs, person_a, person_b): if person_a and person_b: first = next((run for run in runs if person_a.lower() in run["person"]["name"].lower()), None) second = next((run for run in runs if person_b.lower() in run["person"]["name"].lower()), None) if first and second and first["_dim"] == second["_dim"]: return first, second for index, first in enumerate(runs): for second in runs[index + 1 :]: if first["person"]["name"] != second["person"]["name"] and first["_dim"] == second["_dim"]: return first, second raise RuntimeError("No two compatible saved persona runs found") def cosine_distance(a, b): denom = max(float(np.linalg.norm(a) * np.linalg.norm(b)), 1e-12) return 1.0 - float(np.dot(a, b) / denom) def pca_2d(vectors): centered = vectors - vectors.mean(axis=0, keepdims=True) _, _, vt = np.linalg.svd(centered, full_matrices=False) return centered @ vt[:2].T def paired_points(run_a, run_b): vectors = [] labels = [] distances = [] for answer_a in run_a["answers"]: answer_b = next((item for item in run_b["answers"] if item["task_id"] == answer_a["task_id"]), None) if answer_b is None: continue vector_a = np.asarray(answer_a["embedding"], dtype=np.float32) vector_b = np.asarray(answer_b["embedding"], dtype=np.float32) if vector_a.shape != vector_b.shape: continue distance = cosine_distance(vector_a, vector_b) distances.append((answer_a["task_id"], answer_a["category"], distance)) vectors.extend([vector_a, vector_b]) labels.extend( [ {"side": "A", "task": answer_a["task_id"], "category": answer_a["category"]}, {"side": "B", "task": answer_b["task_id"], "category": answer_b["category"]}, ] ) if not distances: raise RuntimeError("Selected runs do not share compatible task embeddings") return np.vstack(vectors), labels, distances def add_margin(values, pad=0.18): low = float(np.min(values)) high = float(np.max(values)) span = high - low if span <= 1e-9: return low - 1.0, high + 1.0 return low - span * pad, high + span * pad def draw(run_a, run_b, output_path): vectors, labels, distances = paired_points(run_a, run_b) coords = pca_2d(vectors) name_a = run_a["person"]["name"] name_b = run_b["person"]["name"] fig = plt.figure(figsize=(15, 8.5), dpi=160, facecolor="#f4f6fb") grid = fig.add_gridspec(1, 2, width_ratios=[1.55, 1], wspace=0.25) ax = fig.add_subplot(grid[0, 0], facecolor="#ffffff") bar = fig.add_subplot(grid[0, 1], facecolor="#ffffff") line_color = "#9ca3af" blue = "#2563eb" red = "#dc2626" for task in sorted({item["task"] for item in labels}): indices = [index for index, item in enumerate(labels) if item["task"] == task] if len(indices) != 2: continue ax.plot(coords[indices, 0], coords[indices, 1], color=line_color, linewidth=1.1, alpha=0.8, zorder=1) for side, color, marker, name, offset in [ ("A", blue, "o", name_a, (7, 7)), ("B", red, "D", name_b, (7, -13)), ]: indices = [index for index, item in enumerate(labels) if item["side"] == side] ax.scatter(coords[indices, 0], coords[indices, 1], s=120, c=color, marker=marker, edgecolors="#ffffff", linewidths=1.5, label=name, zorder=3) for index in indices: ax.annotate(labels[index]["task"], coords[index], xytext=offset, textcoords="offset points", fontsize=8.5, color="#111827") ax.axhline(0, color="#cbd5e1", linewidth=1) ax.axvline(0, color="#cbd5e1", linewidth=1) ax.grid(True, color="#e5e7eb", linewidth=0.8) ax.set_xlim(*add_margin(coords[:, 0])) ax.set_ylim(*add_margin(coords[:, 1])) ax.set_xlabel("PCA axis 1", color="#111827") ax.set_ylabel("PCA axis 2", color="#111827") ax.set_title("Saved answer embeddings projected to 2D", loc="left", fontsize=15, color="#111827", pad=14) ax.legend(loc="upper left", frameon=False) sorted_distances = sorted(distances, key=lambda item: item[2]) y = np.arange(len(sorted_distances)) values = [item[2] for item in sorted_distances] labels_y = [f"{item[0]} ({item[1]})" for item in sorted_distances] bar.barh(y, values, color="#334155", alpha=0.88) bar.set_yticks(y) bar.set_yticklabels(labels_y, fontsize=9, color="#111827") bar.set_xlabel("Cosine distance", color="#111827") bar.set_title("Per-task answer divergence", loc="left", fontsize=15, color="#111827", pad=14) bar.grid(True, axis="x", color="#e5e7eb", linewidth=0.8) for index, value in enumerate(values): bar.text(value + 0.004, index, f"{value:.3f}", va="center", fontsize=8.5, color="#111827") mean_distance = float(np.mean(values)) fig.suptitle(f"{name_a} vs {name_b} | mean cosine distance {mean_distance:.3f}", x=0.06, ha="left", fontsize=18, color="#111827", fontweight="bold") fig.text(0.06, 0.925, f"{run_a['models'].get('embedding_source')} | {len(values)} matched benchmark tasks", ha="left", fontsize=10, color="#475569") output_path.parent.mkdir(parents=True, exist_ok=True) fig.savefig(output_path, bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close(fig) def main(): parser = argparse.ArgumentParser() parser.add_argument("--person-a", default="") parser.add_argument("--person-b", default="") parser.add_argument("--output", default=str(OUTPUT_PATH)) args = parser.parse_args() run_a, run_b = pick_runs(load_runs(), args.person_a, args.person_b) output_path = Path(args.output) draw(run_a, run_b, output_path) print(output_path.resolve()) print(run_a["person"]["name"], run_b["person"]["name"]) if __name__ == "__main__": main()