embedding_transform_eval.py

#!/usr/bin/env python3
"""复用 batch_top100_match.py 缓存的 embedding，试几种空间变换，看哪个最能提分。

方法：
  raw         - 不变换（baseline，用 cosine sim → Top-K 投票）
  severity_axis - 单方向投影：方向 = severe ruler 均值 − 非 severe ruler 均值
  ridge       - 岭回归：ruler_emb → ruler_score，拟合 (w, b)，用 w·emb+b 预测
  lasso       - L1 回归
  lda         - Fisher LDA：把 ruler 二值化（rank<106 vs >=106）找投影方向
  pca128_ridge - 先 PCA 到 128 维再 ridge
  knn_score   - kNN 回归（用 ruler 100 邻居均 score 当预测，本质和 batch_top100 等价）

要求：
  - cache_emb/csv_*.npy + cache_emb/ruler_*.npy（之前跑过 batch_top100_match.py 自动缓存）
  - 标签 csv（拿 golden_set.csv 的 label 列做 GT）
  - ruler_items.json（拿 score / rank）

用法：
  python3 embedding_transform_eval.py
  python3 embedding_transform_eval.py --pca-dim 256 --boundary-rank 106
"""
import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd

from sklearn.linear_model import Ridge, Lasso
from sklearn.decomposition import PCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA


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",
    pos_label = "Y",
    boundary_rank = 106,
)


def load_npy_pair(cache_dir, n_csv, n_ruler, max_length=4096):
    """根据 batch_top100_match.py 的命名规则找缓存。"""
    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 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_f1(scores, gts):
    """扫所有可能阈值，找最大化 F1 的那个。返回 (f1, thr, p, r)。"""
    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 fit_severity_axis(emb, ruler_score, ruler_rank, boundary_rank):
    """方向 = 严重组均值 - 非严重组均值；投影 = emb @ direction。"""
    severe = emb[ruler_rank < boundary_rank].mean(axis=0)
    notsev = emb[ruler_rank >= boundary_rank].mean(axis=0)
    direction = severe - notsev
    direction = direction / (np.linalg.norm(direction) + 1e-12)
    return direction


def project(emb, direction):
    return emb @ direction


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=DEFAULTS["pos_label"])
    p.add_argument("--boundary-rank", type=int, default=DEFAULTS["boundary_rank"])
    p.add_argument("--pca-dim",       type=int, default=128)
    p.add_argument("--max-length",    type=int, default=4096)
    p.add_argument("--output-jsonl",  default="transform_eval.jsonl")
    args = p.parse_args()

    print("[1] load labels and ruler meta")
    df = pd.read_csv(args.csv, keep_default_na=False)
    gts = (df[df.columns[df.columns.tolist().index("label")]]
           .astype(str).str.upper().eq(args.positive_label.upper()).astype(int).values)
    ruler_rank, ruler_score = load_ruler_meta(args.ruler)
    n_csv = len(gts)
    n_ruler = len(ruler_rank)
    print(f"    csv={n_csv}, ruler={n_ruler}, pos rate={gts.mean():.2%}")

    print("[2] load embeddings from cache")
    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}")
    # 已经是 L2 归一化的（前面脚本里做了）

    methods = {}

    # --- raw cosine top-K weighted ---
    K = 100
    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_scores = ruler_score[top_idx]
    raw_weighted = (top_sims * top_scores).sum(axis=1) / np.maximum(top_sims.sum(axis=1), 1e-12)
    methods["raw cosine + top100 weighted score"] = raw_weighted

    # --- severity axis projection ---
    direction = fit_severity_axis(ruler_emb, ruler_score, ruler_rank, args.boundary_rank)
    methods["severity_axis projection (1D)"] = project(csv_emb, direction)

    # --- ridge regression: emb -> score ---
    rid = Ridge(alpha=1.0).fit(ruler_emb, ruler_score)
    methods["ridge: emb -> score"] = rid.predict(csv_emb)

    # --- lasso ---
    las = Lasso(alpha=0.001, max_iter=5000).fit(ruler_emb, ruler_score)
    methods["lasso: emb -> score"] = las.predict(csv_emb)

    # --- LDA: severe/notsevere ---
    y_bin = (ruler_rank < args.boundary_rank).astype(int)
    lda = LDA().fit(ruler_emb, y_bin)
    methods["LDA: severe vs not"] = lda.decision_function(csv_emb)

    # --- PCA -> ridge ---
    pca = PCA(n_components=min(args.pca_dim, n_ruler-1, ruler_emb.shape[1])).fit(ruler_emb)
    rid_p = Ridge(alpha=1.0).fit(pca.transform(ruler_emb), ruler_score)
    methods[f"PCA{pca.n_components_} + ridge"] = rid_p.predict(pca.transform(csv_emb))

    # --- knn average top-100 ruler score ---
    methods["kNN-100 mean(ruler_score)"] = top_scores.mean(axis=1)

    # --- LLM 列（如果 csv 里带了 AIPF 跑出来的位置/score）---
    BOUNDARY_SCORE_DEFAULT = 44.72
    llm_cols = [
        ("score_gemini_2.5_flash",     None),       # 已经是 score，越大越严
        ("position_gemini_2.5_flash",  "neg"),      # position 越小越严，取负
        ("score_gpt_4.1",              None),
        ("position_gpt_4.1",           "neg"),
    ]
    for col, mode in llm_cols:
        if col not in df.columns:
            continue
        raw = pd.to_numeric(df[col], errors="coerce").values
        # NaN 用列中位数填，避免阈值扫描出问题
        med = np.nanmedian(raw)
        if np.isnan(med):
            continue
        raw = np.where(np.isnan(raw), med, raw)
        if mode == "neg":
            methods[f"LLM: {col} (-position)"] = -raw
        else:
            methods[f"LLM: {col}"] = raw

    # ---- 评分输出 ----
    print(f"\n{'method':<40}{'best F1':>10}{'thr':>10}{'P':>9}{'R':>9}{'AUC?':>10}")
    print("-" * 88)
    rows = []
    for name, scores in methods.items():
        f1, thr, prec, rec = best_threshold_f1(scores, gts)
        try:
            from sklearn.metrics import roc_auc_score
            auc = roc_auc_score(gts, scores)
        except Exception:
            auc = float("nan")
        rows.append((name, f1, thr, prec, rec, auc))
        print(f"{name:<40}{f1:>10.4f}{thr:>10.4f}{prec:>9.4f}{rec:>9.4f}{auc:>10.4f}")

    # 写 jsonl 把每条样本 7 个分数都留下
    print(f"\n[write] {args.output_jsonl}")
    with open(args.output_jsonl, "w") as f:
        for i in range(n_csv):
            rec = {"i": i, "ground_truth": int(gts[i])}
            for name, scores in methods.items():
                rec[name] = float(scores[i])
            f.write(json.dumps(rec) + "\n")

    print("\n说明：")
    print("- AUC 反映分布可分性，跟阈值无关。AUC 高 = 这个变换的输出能更好把正/负分开。")
    print("- best F1 是扫阈值找到的上限，是这个变换的理论最佳。")


if __name__ == "__main__":
    main()