add script: write top-100 estimated position back to golden_set

add_estimated_position.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""把 top100 match 估出来的"均值 position" 写回 golden_set.csv。
+
+逻辑：
+  - 每条 golden 样本（共 1000 条）已经在 golden_top100.jsonl 里有 100 个最近邻 ruler item
+  - 对这 100 个邻居的 rank 求均值（也可以 sim-加权），作为 embedding 估出的 position
+  - 通过 task_id 把这个值合并回 csv，新列：estimated_position_top100
+  - 同时也保留几个有用的衍生字段：weighted_position / median_position / mean_score
+
+输出：
+  - golden_set_with_estimated.csv （默认；不覆盖原文件）
+
+用法：
+  python3 add_estimated_position.py
+  python3 add_estimated_position.py \
+    --csv /mnt/.../aipf_golden_set.csv \
+    --jsonl golden_top100.jsonl \
+    --output golden_set_with_estimated.csv
+"""
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+
+DEFAULTS = dict(
+    csv    = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/example/yss_ruler_eval/data/aipf_golden_set.csv",
+    jsonl  = "golden_top100.jsonl",
+    output = "golden_set_with_estimated.csv",
+    id_col = "task_id",
+)
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--csv",    default=DEFAULTS["csv"])
+    p.add_argument("--jsonl",  default=DEFAULTS["jsonl"])
+    p.add_argument("--output", default=DEFAULTS["output"])
+    p.add_argument("--id-col", default=DEFAULTS["id_col"])
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # 1) 读 jsonl，按 task_id 算各种聚合
+    print(f"[1] read jsonl: {args.jsonl}")
+    estimated = {}
+    with open(args.jsonl, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            r = json.loads(line)
+            ranks  = np.array([t["rank"]  for t in r["topk"]], dtype=float)
+            sims   = np.array([t["sim"]   for t in r["topk"]], dtype=float)
+            scores = np.array([t["score"] for t in r["topk"]], dtype=float)
+            wsim = sims.sum()
+
+            estimated[str(r["task_id"])] = {
+                "estimated_position_top100":          float(ranks.mean()),
+                "estimated_position_top100_weighted": float((sims * ranks).sum() / max(wsim, 1e-12)),
+                "estimated_position_top100_median":   float(np.median(ranks)),
+                "estimated_score_top100_mean":        float(scores.mean()),
+                "estimated_score_top100_weighted":    float((sims * scores).sum() / max(wsim, 1e-12)),
+                "top1_rank": int(ranks[0]),
+                "top1_sim":  float(sims[0]),
+            }
+    print(f"      -> {len(estimated)} rows in jsonl")
+
+    # 2) 读 csv
+    print(f"[2] read csv: {args.csv}")
+    df = pd.read_csv(args.csv, keep_default_na=False)
+    print(f"      -> {len(df)} rows, {len(df.columns)} cols")
+    print(f"      columns: {list(df.columns)}")
+
+    if args.id_col not in df.columns:
+        raise ValueError(f"id column '{args.id_col}' not in csv. available: {list(df.columns)}")
+
+    # 3) merge
+    print(f"[3] merge by {args.id_col}")
+    new_cols = [
+        "estimated_position_top100",
+        "estimated_position_top100_weighted",
+        "estimated_position_top100_median",
+        "estimated_score_top100_mean",
+        "estimated_score_top100_weighted",
+        "top1_rank",
+        "top1_sim",
+    ]
+    for col in new_cols:
+        df[col] = df[args.id_col].astype(str).map(lambda k: estimated.get(k, {}).get(col))
+
+    matched = df["estimated_position_top100"].notna().sum()
+    missing = len(df) - matched
+    print(f"      matched: {matched}/{len(df)}")
+    if missing:
+        print(f"      ⚠️ missing {missing} task_id（jsonl 里没有这些 id）")
+
+    # 4) 写出
+    print(f"[4] write: {args.output}")
+    df.to_csv(args.output, index=False)
+    print(f"      done")
+
+    # 5) sanity: 看几行
+    print("\n[preview] 前 3 行新列：")
+    cols_show = [args.id_col, "label"] + new_cols
+    cols_show = [c for c in cols_show if c in df.columns]
+    print(df[cols_show].head(3).to_string(index=False))
+
+    # 6) 简单 stats
+    print("\n[stats] estimated_position_top100 分布:")
+    s = pd.to_numeric(df["estimated_position_top100"], errors="coerce")
+    print(f"  min={s.min():.2f}  max={s.max():.2f}  mean={s.mean():.2f}  median={s.median():.2f}")
+
+    if "label" in df.columns:
+        print("\n[stats] 按 label 分组:")
+        for lab, sub in df.groupby("label"):
+            ss = pd.to_numeric(sub["estimated_position_top100"], errors="coerce")
+            print(f"  label={lab!r:>4}  n={len(sub):>4}  mean={ss.mean():.2f}  median={ss.median():.2f}")
+
+
+if __name__ == "__main__":
+    main()