support arbitrary K for estimated_position (top1/top5/top100)

add_estimated_position.py

@@ -1,16 +1,25 @@
 #!/usr/bin/env python3
-"""把 top100 match 估出来的整数 position 写回 golden_set.csv。
+"""把 embedding top-K match 估出来的整数 position 写回 golden_set.csv。
 
-每条 case：
-  - 从 golden_top100.jsonl 取它的 100 个最近邻 ruler 的 rank
-  - 求均值，round 成整数
-  - 作为新列 estimated_position 写回 csv
-
-输出：golden_set_with_estimated.csv（不覆盖原文件）
+支持任意 K（默认 100），列名自动带 K 后缀，方便同时跑多种 K 比较。
 
 用法：
-  python3 add_estimated_position.py
-  python3 add_estimated_position.py --csv /path/to.csv --jsonl /path/to.jsonl
+  # top-100 邻居均值
+  python3 add_estimated_position.py --k 100
+
+  # top-5 邻居均值
+  python3 add_estimated_position.py --k 5
+
+  # top-1（最近邻 rank）
+  python3 add_estimated_position.py --k 1
+
+  # 一次跑多种 K（每种生成一列）
+  python3 add_estimated_position.py --k 1 --k-suffix top1
+  python3 add_estimated_position.py --csv golden_set_with_estimated.csv --k 5 --k-suffix top5
+  python3 add_estimated_position.py --csv golden_set_with_estimated.csv --k 100 --k-suffix top100
+
+  # 自定义列名
+  python3 add_estimated_position.py --k 5 --new-col estimated_position_top5
 """
 import argparse
 import json
@@ -25,78 +34,92 @@ DEFAULTS = dict(
     jsonl  = "golden_top100.jsonl",
     output = "golden_set_with_estimated.csv",
     id_col = "task_id",
-    new_col = "estimated_position",
 )
 
 
 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"])
-    p.add_argument("--new-col", default=DEFAULTS["new_col"])
+    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"])
+    p.add_argument("--k",        type=int, default=100,
+                   help="用前 K 个邻居算均值（K=1 即 top1 rank）")
+    p.add_argument("--new-col",  default=None,
+                   help="新列名；不传则自动用 estimated_position_top{K}")
+    p.add_argument("--k-suffix", default=None,
+                   help="（已弃用，等价于 --new-col estimated_position_{suffix}）")
     return p.parse_args()
 
 
 def main():
     args = parse_args()
+    K = args.k
 
-    # 1) 读 jsonl，每条样本算 top-100 rank 均值 → 整数
+    if args.new_col:
+        col_name = args.new_col
+    elif args.k_suffix:
+        col_name = f"estimated_position_{args.k_suffix}"
+    else:
+        col_name = f"estimated_position_top{K}"
+
+    # 1) 读 jsonl，每条样本算 top-K rank 均值（K=1 时等价于最近邻 rank）
     print(f"[1] read jsonl: {args.jsonl}")
     estimated = {}
+    skipped = 0
     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)
+            topk = r.get("topk") or []
+            if len(topk) < K:
+                skipped += 1
+                continue
+            ranks = np.array([t["rank"] for t in topk[:K]], dtype=float)
             estimated[str(r["task_id"])] = int(round(ranks.mean()))
-    print(f"      -> {len(estimated)} cases estimated")
+    print(f"      -> {len(estimated)} cases estimated (K={K})")
+    if skipped:
+        print(f"      ⚠️ {skipped} 条 jsonl 邻居数 < K，跳过")
 
     # 2) 读 csv
     print(f"[2] read csv: {args.csv}")
     df = pd.read_csv(args.csv, keep_default_na=False)
-    print(f"      -> {len(df)} rows")
+    print(f"      -> {len(df)} rows, {len(df.columns)} cols")
 
     if args.id_col not in df.columns:
         raise ValueError(f"'{args.id_col}' not in csv. cols: {list(df.columns)}")
 
-    # 3) 给每条 case 加 estimated_position 列
-    print(f"[3] add column: {args.new_col}")
-    df[args.new_col] = df[args.id_col].astype(str).map(estimated)
+    # 3) 给每条 case 加新列
+    print(f"[3] add column: {col_name}")
+    df[col_name] = df[args.id_col].astype(str).map(estimated)
 
-    n_total   = len(df)
-    n_matched = df[args.new_col].notna().sum()
-    n_missing = n_total - n_matched
+    n_total = len(df)
+    n_matched = df[col_name].notna().sum()
     print(f"      matched: {n_matched}/{n_total}")
-    if n_missing:
-        print(f"      ⚠️ {n_missing} 条 case 在 jsonl 里没找到对应 task_id")
-    else:
-        # 全部匹配上 → 转成 Int64（pandas 的可空整数类型，不会因为没 NaN 变 float）
-        df[args.new_col] = df[args.new_col].astype("Int64")
+    if n_matched == n_total:
+        df[col_name] = df[col_name].astype("Int64")
 
     # 4) 写
     print(f"[4] write: {args.output}")
     df.to_csv(args.output, index=False)
-    print(f"      done. csv 现在有 {len(df.columns)} 列（多了 1 列：{args.new_col}）")
+    print(f"      done. csv 现在有 {len(df.columns)} 列（含新列：{col_name}）")
 
-    # 5) 预览
-    print(f"\n[preview] 前 5 行:")
-    cols = [args.id_col, "label", args.new_col]
+    # 5) 预览 + stats
+    print(f"\n[preview] 前 5 行（task_id / label / {col_name}）:")
+    cols = [args.id_col, "label", col_name]
     cols = [c for c in cols if c in df.columns]
     print(df[cols].head(5).to_string(index=False))
 
-    # 6) 简单 stats
-    s = pd.to_numeric(df[args.new_col], errors="coerce")
-    print(f"\n[stats] {args.new_col}:")
-    print(f"  min={int(s.min())}  max={int(s.max())}  mean={s.mean():.1f}  median={int(s.median())}")
+    s = pd.to_numeric(df[col_name], errors="coerce")
+    print(f"\n[stats] {col_name}:")
+    print(f"  min={int(s.min())}  max={int(s.max())}  mean={s.mean():.2f}  median={int(s.median())}  std={s.std():.2f}")
     if "label" in df.columns:
         print("\n按 label 分组:")
         for lab, sub in df.groupby("label"):
-            ss = pd.to_numeric(sub[args.new_col], errors="coerce")
-            print(f"  label={lab!r:>4}  n={len(sub):>4}  mean={ss.mean():.1f}  median={int(ss.median())}")
+            ss = pd.to_numeric(sub[col_name], errors="coerce")
+            print(f"  label={lab!r:>4}  n={len(sub):>4}  mean={ss.mean():.2f}  median={int(ss.median())}  std={ss.std():.2f}")
 
 
 if __name__ == "__main__":