add full t-SNE: 1000 golden + 200 ruler

full_tsne.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+"""把 1000 golden + 200 ruler 一起做 t-SNE 可视化。
+
+复用 cache_emb/ 里缓存的 embedding，不重算。
+默认 1200 个点全画出来，golden 按 GT label 着色，ruler 按 score 着色。
+
+用法：
+  python3 full_tsne.py
+  python3 full_tsne.py --perplexity 30 --output full_tsne.png
+  python3 full_tsne.py --label-rulers boundary   # 只标 boundary 附近的 ruler rank
+  python3 full_tsne.py --label-rulers all        # 标全部 200 个 ruler rank（密）
+  python3 full_tsne.py --label-rulers none       # 不标 rank（最干净）
+"""
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+
+
+DEFAULTS = dict(
+    cache_dir = "cache_emb",
+    csv       = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/example/yss_ruler_eval/data/aipf_golden_set.csv",
+    ruler     = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/ranking_moderation/data/dm/youth_sexual_and_physical_abuse_aigt_v009/ranking_bucket/ruler_items.json",
+)
+
+
+def load_npy_pair(cache_dir, n_csv, n_ruler, max_length=4096):
+    cd = Path(cache_dir)
+    csvs   = list(cd.glob(f"csv_*_n{n_csv}_L{max_length}.npy"))
+    rulers = list(cd.glob(f"ruler_*_n{n_ruler}_L{max_length}.npy"))
+    if not csvs or not rulers:
+        raise FileNotFoundError(f"找不到缓存。期望 {cd}/csv_*_n{n_csv}_L{max_length}.npy 和 ruler_*_n{n_ruler}_L{max_length}.npy")
+    return np.load(csvs[0]), np.load(rulers[0])
+
+
+def load_ruler_meta(path):
+    with open(path) as f:
+        data = json.load(f)
+    items = data if isinstance(data, list) else (data.get("items") or data.get("ruler_items") or data.get("data") or [])
+    ranks  = np.array([int(it["rank"])  for it in items])
+    scores = np.array([float(it["score"]) for it in items])
+    return ranks, scores
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--cache-dir",     default=DEFAULTS["cache_dir"])
+    p.add_argument("--csv",           default=DEFAULTS["csv"])
+    p.add_argument("--ruler",         default=DEFAULTS["ruler"])
+    p.add_argument("--positive-label", default="Y")
+    p.add_argument("--boundary-rank", type=int, default=106)
+    p.add_argument("--max-length",    type=int, default=4096)
+    p.add_argument("--perplexity",    type=float, default=30.0)
+    p.add_argument("--output",        default="full_tsne.png")
+    p.add_argument("--label-rulers",  default="boundary",
+                   choices=["all", "boundary", "none"],
+                   help="all=全部标; boundary=只标 boundary±10 + 端点; none=不标")
+    p.add_argument("--label-fontsize", type=float, default=5)
+    return run(p.parse_args())
+
+
+def run(args):
+    print("[1] load labels and ruler meta")
+    df = pd.read_csv(args.csv, keep_default_na=False)
+    gts = df["label"].astype(str).str.upper().eq(args.positive_label.upper()).astype(int).values
+    ruler_rank, ruler_score = load_ruler_meta(args.ruler)
+    n_csv, n_ruler = len(gts), len(ruler_rank)
+    print(f"    csv={n_csv}, ruler={n_ruler}")
+
+    print("[2] load cached embeddings")
+    csv_emb, ruler_emb = load_npy_pair(args.cache_dir, n_csv, n_ruler, args.max_length)
+    print(f"    csv_emb={csv_emb.shape}, ruler_emb={ruler_emb.shape}")
+
+    all_emb = np.vstack([csv_emb, ruler_emb])
+    print(f"[3] t-SNE on {all_emb.shape[0]} points (perplexity={args.perplexity})")
+    tsne = TSNE(n_components=2, perplexity=args.perplexity,
+                init="pca", random_state=42, metric="cosine",
+                learning_rate="auto")
+    xy = tsne.fit_transform(all_emb)
+    csv_xy = xy[:n_csv]
+    ruler_xy = xy[n_csv:]
+
+    print(f"[4] draw to {args.output}")
+    fig, ax = plt.subplots(figsize=(16, 13), dpi=130)
+
+    # golden A（GT=0）灰色淡
+    a_mask = gts == 0
+    ax.scatter(csv_xy[a_mask, 0], csv_xy[a_mask, 1],
+               c="#bbbbbb", s=12, alpha=0.55,
+               edgecolor="none", label=f"golden A (N={a_mask.sum()})")
+
+    # golden Y（GT=1）红色
+    y_mask = gts == 1
+    ax.scatter(csv_xy[y_mask, 0], csv_xy[y_mask, 1],
+               c="#d62728", s=14, alpha=0.85,
+               edgecolor="none", label=f"golden Y (N={y_mask.sum()})")
+
+    # ruler 200 用 colormap，按 score 着色，加黑边圈以区分
+    sc = ax.scatter(ruler_xy[:, 0], ruler_xy[:, 1],
+                    c=ruler_score, cmap="RdYlGn_r",
+                    s=70, alpha=0.95,
+                    edgecolor="black", linewidth=0.8,
+                    label=f"ruler (N={n_ruler})")
+    cbar = plt.colorbar(sc, ax=ax, shrink=0.7)
+    cbar.set_label("ruler score (high = severe)")
+
+    # 标 ruler rank
+    if args.label_rulers == "all":
+        for (x, y), r in zip(ruler_xy, ruler_rank):
+            ax.annotate(str(r), (x, y), fontsize=args.label_fontsize,
+                        ha="center", va="center", alpha=0.9)
+    elif args.label_rulers == "boundary":
+        # 标端点 + boundary 附近
+        targets = set([0, n_ruler - 1])
+        targets.update(range(max(0, args.boundary_rank - 10),
+                             min(n_ruler, args.boundary_rank + 11)))
+        for i, ((x, y), r) in enumerate(zip(ruler_xy, ruler_rank)):
+            if int(r) in targets:
+                ax.annotate(str(r), (x, y),
+                            fontsize=args.label_fontsize + 2,
+                            ha="center", va="center",
+                            fontweight="bold")
+
+    ax.set_title(f"t-SNE: 1000 golden + 200 ruler (Qwen3-Embedding-8B)\n"
+                 f"red=GT=Y, gray=GT=A, color circles=ruler items "
+                 f"(rank={'/'.join(['0', str(args.boundary_rank), str(n_ruler-1)])} highlighted)")
+    ax.set_xlabel("t-SNE 1")
+    ax.set_ylabel("t-SNE 2")
+    ax.legend(loc="best", framealpha=0.9)
+    plt.tight_layout()
+    plt.savefig(args.output, dpi=130, bbox_inches="tight")
+    print(f"    saved: {args.output}")
+
+    # 顺便算一个简单的"几何分离度"
+    y_centroid = csv_xy[y_mask].mean(axis=0)
+    a_centroid = csv_xy[a_mask].mean(axis=0)
+    ruler_top = ruler_xy[ruler_rank < args.boundary_rank].mean(axis=0)
+    ruler_bot = ruler_xy[ruler_rank >= args.boundary_rank].mean(axis=0)
+    print("\n[centroids in t-SNE space]")
+    print(f"  golden Y     -> {y_centroid}")
+    print(f"  golden A     -> {a_centroid}")
+    print(f"  ruler severe -> {ruler_top}")
+    print(f"  ruler not    -> {ruler_bot}")
+    print(f"  Y-A distance     : {np.linalg.norm(y_centroid - a_centroid):.2f}")
+    print(f"  Y-ruler-severe   : {np.linalg.norm(y_centroid - ruler_top):.2f}")
+    print(f"  A-ruler-not      : {np.linalg.norm(a_centroid - ruler_bot):.2f}")
+
+
+if __name__ == "__main__":
+    main()