| |
| """把 embedding top-K match 估出来的整数 position 写回 golden_set.csv。 |
| |
| 默认列名固定为 estimated_position（下游 AIPF warm-start 代码读这个名）。 |
| 默认输入/输出都指向 AIPF 流水线实际读取的那个 csv，**直接原地覆盖**， |
| 省掉 `cp` 到 example/yss_ruler_eval/data/ 这一步。 |
| |
| 每次跑都覆盖 estimated_position 列；用 --k 切换不同邻居数即可， |
| 其它列保持不变。 |
| |
| 用法： |
| python3 add_estimated_position.py --k 100 # 默认 K |
| python3 add_estimated_position.py --k 5 |
| python3 add_estimated_position.py --k 1 |
| |
| ⚠️ 第一次跑前建议备份原 csv： |
| cp /mnt/.../aipf_golden_set.csv /mnt/.../aipf_golden_set.csv.bak |
| |
| 可选：想同时存多列对比（不覆盖默认列），手动指定列名： |
| python3 add_estimated_position.py --k 5 --new-col estimated_position_top5 |
| """ |
| 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 = "/mnt/bn/tns-algo-ue-my/biaowu/aipf_dm_metric/example/yss_ruler_eval/data/aipf_golden_set.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"]) |
| 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 |
|
|
| |
| |
| if args.new_col: |
| col_name = args.new_col |
| elif args.k_suffix: |
| col_name = f"estimated_position_{args.k_suffix}" |
| else: |
| col_name = "estimated_position" |
|
|
| |
| 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) |
| 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 (K={K})") |
| if skipped: |
| print(f" ⚠️ {skipped} 条 jsonl 邻居数 < K，跳过") |
|
|
| |
| 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") |
|
|
| if args.id_col not in df.columns: |
| raise ValueError(f"'{args.id_col}' not in csv. cols: {list(df.columns)}") |
|
|
| |
| 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[col_name].notna().sum() |
| print(f" matched: {n_matched}/{n_total}") |
| if n_matched == n_total: |
| df[col_name] = df[col_name].astype("Int64") |
|
|
| |
| print(f"[4] write: {args.output}") |
| df.to_csv(args.output, index=False) |
| print(f" done. csv 现在有 {len(df.columns)} 列（含新列：{col_name}）") |
|
|
| |
| 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)) |
|
|
| 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[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__": |
| main() |
|
|
