simplify: only add one estimated_position column (int)

add_estimated_position.py

@@ -1,21 +1,16 @@
 #!/usr/bin/env python3
-"""把 top100 match 估出来的"均值 position" 写回 golden_set.csv。
+"""把 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
+每条 case：
+  - 从 golden_top100.jsonl 取它的 100 个最近邻 ruler 的 rank
+  - 求均值，round 成整数
+  - 作为新列 estimated_position 写回 csv
 
-输出：
-  - golden_set_with_estimated.csv （默认；不覆盖原文件）
+输出：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
+  python3 add_estimated_position.py --csv /path/to.csv --jsonl /path/to.jsonl
 """
 import argparse
 import json
@@ -30,22 +25,24 @@ 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("--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"])
     return p.parse_args()
 
 
 def main():
     args = parse_args()
 
-    # 1) 读 jsonl，按 task_id 算各种聚合
+    # 1) 读 jsonl，每条样本算 top-100 rank 均值 → 整数
     print(f"[1] read jsonl: {args.jsonl}")
     estimated = {}
     with open(args.jsonl, "r", encoding="utf-8") as f:
@@ -54,77 +51,52 @@ def main():
             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()
-
-            mean_pos     = ranks.mean()
-            weighted_pos = (sims * ranks).sum() / max(wsim, 1e-12)
-            median_pos   = np.median(ranks)
-            estimated[str(r["task_id"])] = {
-                # position 类一律 round 成整数（和 AIPF 的 ruler_position 对齐）
-                "estimated_position_top100":          int(round(mean_pos)),
-                "estimated_position_top100_weighted": int(round(weighted_pos)),
-                "estimated_position_top100_median":   int(round(median_pos)),
-                # score 保留浮点
-                "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")
+            ranks = np.array([t["rank"] for t in r["topk"]], dtype=float)
+            estimated[str(r["task_id"])] = int(round(ranks.mean()))
+    print(f"      -> {len(estimated)} cases estimated")
 
     # 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)}")
+    print(f"      -> {len(df)} rows")
 
     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) 写出
+        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)
+
+    n_total   = len(df)
+    n_matched = df[args.new_col].notna().sum()
+    n_missing = n_total - n_matched
+    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")
+
+    # 4) 写
     print(f"[4] write: {args.output}")
     df.to_csv(args.output, index=False)
-    print(f"      done")
+    print(f"      done. csv 现在有 {len(df.columns)} 列（多了 1 列：{args.new_col}）")
 
-    # 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))
+    # 5) 预览
+    print(f"\n[preview] 前 5 行:")
+    cols = [args.id_col, "label", args.new_col]
+    cols = [c for c in cols if c in df.columns]
+    print(df[cols].head(5).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}")
-
+    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())}")
     if "label" in df.columns:
-        print("\n[stats] 按 label 分组:")
+        print("\n按 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}")
+            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())}")
 
 
 if __name__ == "__main__":