scripts/export_embedding_comparison.py

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()