compare 4096D cosine top-100 vs t-SNE 2D top-100

tsne_circle_eval.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+"""对比：4096D cosine 邻居 vs t-SNE 2D 圆内邻居，看哪种判别更准。
+
+复用 cache_emb/。t-SNE 用 1000 golden + 200 ruler 一起做（1200 点），保证投影一致。
+"""
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+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")
+    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 metrics(preds, gts):
+    tp = int(((preds == 1) & (gts == 1)).sum())
+    fp = int(((preds == 1) & (gts == 0)).sum())
+    tn = int(((preds == 0) & (gts == 0)).sum())
+    fn = int(((preds == 0) & (gts == 1)).sum())
+    p = tp/(tp+fp) if tp+fp else 0.0
+    r = tp/(tp+fn) if tp+fn else 0.0
+    f = 2*p*r/(p+r) if p+r else 0.0
+    a = (tp+tn)/len(preds)
+    return tp, fp, tn, fn, p, r, f, a
+
+
+def best_threshold(scores, gts):
+    cands = sorted(set(scores.tolist()))
+    best = (-1.0, None, None, None)
+    for c in cands:
+        preds = (scores >= c).astype(int)
+        _, _, _, _, p, r, f, _ = metrics(preds, gts)
+        if f > best[0]:
+            best = (f, c, p, r)
+    return best
+
+
+def topk_neighbors(query_xy, ruler_xy, k):
+    """对每个 query，找 ruler 里最近的 k 个，返回 (idx, dist)"""
+    # query_xy (Nq, 2), ruler_xy (Nr, 2)
+    diffs = query_xy[:, None, :] - ruler_xy[None, :, :]
+    dists = np.linalg.norm(diffs, axis=-1)              # (Nq, Nr)
+    idx = np.argpartition(dists, k - 1, axis=1)[:, :k]
+    row = np.arange(len(query_xy))[:, None]
+    selected = dists[row, idx]
+    order = np.argsort(selected, axis=1)
+    return np.take_along_axis(idx, order, axis=1)
+
+
+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("--k",             type=int, default=100)
+    p.add_argument("--seed",          type=int, default=42)
+    args = p.parse_args()
+
+    print("[1] load")
+    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)
+    csv_emb, ruler_emb = load_npy_pair(args.cache_dir, n_csv, n_ruler, args.max_length)
+
+    K = args.k
+    methods = {}
+
+    # ---- baseline: 4096D cosine ----
+    print(f"[2] baseline: 4096D cosine top-{K}")
+    sims = csv_emb @ ruler_emb.T
+    top_idx = np.argpartition(-sims, K-1, axis=1)[:, :K]
+    row = np.arange(n_csv)[:, None]
+    top_sims = sims[row, top_idx]
+    top_score_4096 = ruler_score[top_idx]
+    raw_w = (top_sims * top_score_4096).sum(axis=1) / np.maximum(top_sims.sum(axis=1), 1e-12)
+    raw_mean = top_score_4096.mean(axis=1)
+    raw_vote = (ruler_rank[top_idx] < args.boundary_rank).sum(axis=1)
+    methods["4096D cosine | weighted_score"] = raw_w
+    methods["4096D cosine | mean(score)"]    = raw_mean
+    methods["4096D cosine | vote_count"]     = raw_vote.astype(float)
+
+    # ---- t-SNE 2D ----
+    print(f"[3] t-SNE on 1200 points (perplexity={args.perplexity})")
+    all_emb = np.vstack([csv_emb, ruler_emb])
+    tsne = TSNE(n_components=2, perplexity=args.perplexity,
+                init="pca", random_state=args.seed,
+                metric="cosine", learning_rate="auto")
+    xy = tsne.fit_transform(all_emb)
+    csv_xy, ruler_xy = xy[:n_csv], xy[n_csv:]
+
+    # 2D top-K（等价于"圆扩张到正好包含 100 个 ruler"）
+    print(f"[4] 2D Euclidean top-{K} (in t-SNE space)")
+    top_idx_2d = topk_neighbors(csv_xy, ruler_xy, K)
+    top_score_2d = ruler_score[top_idx_2d]
+    rank_2d = ruler_rank[top_idx_2d]
+    methods["t-SNE 2D | mean(score)"] = top_score_2d.mean(axis=1)
+    methods["t-SNE 2D | vote_count"]  = (rank_2d < args.boundary_rank).sum(axis=1).astype(float)
+    # weighted by 1/dist？也试一下
+    diffs = csv_xy[:, None, :] - ruler_xy[None, :, :]
+    dists2d = np.linalg.norm(diffs, axis=-1)
+    selected_dist = np.take_along_axis(dists2d, top_idx_2d, axis=1)
+    weights = 1.0 / (selected_dist + 1e-6)
+    weighted_2d = (weights * top_score_2d).sum(axis=1) / weights.sum(axis=1)
+    methods["t-SNE 2D | inv_dist weighted"] = weighted_2d
+
+    # ---- 邻居重叠率 ----
+    overlap = []
+    for i in range(n_csv):
+        a = set(top_idx[i].tolist())
+        b = set(top_idx_2d[i].tolist())
+        overlap.append(len(a & b) / K)
+    print(f"\n[5] 邻居重叠率（4096D vs 2D 各取 top-{K}）：")
+    print(f"    平均 = {np.mean(overlap):.2%}")
+    print(f"    中位数 = {np.median(overlap):.2%}")
+    print(f"    p10 / p90 = {np.percentile(overlap, 10):.2%} / {np.percentile(overlap, 90):.2%}")
+
+    # ---- 各方法 best F1 ----
+    print(f"\n[6] best F1 by sweeping threshold (K={K})")
+    print(f"{'method':<35}{'F1':>9}{'thr':>10}{'P':>9}{'R':>9}")
+    print("-" * 75)
+    for name, scores in methods.items():
+        f1, thr, prec, rec = best_threshold(scores, gts)
+        print(f"{name:<35}{f1:>9.4f}{thr:>10.4f}{prec:>9.4f}{rec:>9.4f}")
+
+
+if __name__ == "__main__":
+    main()