tsne_circle_eval.py

#!/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()