chore: import upstream snapshot with attribution

PFPO/scripts/math/rerank_w_orm.py

@@ -0,0 +1,175 @@
+import json
+import argparse
+import os.path
+from glob import glob
+from tqdm import tqdm
+import collections
+from multiprocessing import Pool
+from functools import partial
+import torch
+import sys
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from post_processors.openai_api_callback import majority_voting_predict
+from data.deepseek_math_utils import eval_script
+
+
+def load_rewards(reward_file):
+    if os.path.exists(reward_file):
+        rewards = json.load(open(reward_file))
+    else:
+        rewards = []
+        for file in glob(reward_file):
+            print(file)
+            rewards.extend(json.load(open(file)))
+
+    return rewards
+
+
+def _init(id2reward):
+    global _id2reward
+    _id2reward = id2reward
+
+
+def _worker(item, sc_top_k=None):
+    sorted_results = []
+    if not item["response"] or not item["pred"] or not item["res"]:
+        return {
+            "missing": 1,
+            "reward_missing": 0,
+            "pred_missing": 0,
+            "seq_too_long": 0,
+            "sorted_results": [],
+            "sc_res": False
+        }
+    reward_missing = 0
+    pred_missing = 0
+    seq_too_long = 0
+    for i, (resp, pred, r) in enumerate(zip(item["response"], item["pred"], item["res"])):
+        resp_id = f"{item['id']}_{i}"
+        if resp_id not in _id2reward:
+            reward_missing += 1
+            continue
+        if not pred:
+            pred_missing += 1
+            continue
+
+        reward = _id2reward[resp_id]
+        sorted_results.append((resp, pred, r, reward))
+    sorted_results = sorted(sorted_results, key=lambda x: x[-1], reverse=True)
+
+    sc_top_k_res = {}
+    if sc_top_k and sorted_results:
+        for k in sc_top_k:
+            preds = [r[1] for r in sorted_results[:k]]
+            sc_pred = majority_voting_predict(preds)
+            if sc_pred != "":
+                sc_res = eval_script.eval_math({"prediction": sc_pred, "answer": item["label"]})
+            else:
+                sc_res = False
+            sc_top_k_res[k] = sc_res
+
+    return {
+        "missing": 0,
+        "reward_missing": reward_missing,
+        "pred_missing": pred_missing,
+        "seq_too_long": seq_too_long,
+        "sorted_results": sorted_results,
+        "sc_res": item["sc_res"],
+        "sc_top_k_res": sc_top_k_res,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--response_file", type=str, required=True)
+    parser.add_argument("--reward_file", type=str, required=True)
+    parser.add_argument("--num_workers", type=int, default=8)
+    args = parser.parse_args()
+
+    if os.path.exists(args.response_file):
+        responses = json.load(open(args.response_file))
+    else:
+        responses = []
+        for file in glob(args.response_file):
+            print(file)
+            responses.extend(json.load(open(file)))
+
+    rewards = load_rewards(args.reward_file)
+    id2reward = {item["index"]: item["reward"][1] for item in rewards}
+
+    _k = [1, 3, 5]
+    orm_pass_at_k = {k: 0 for k in _k}
+    missing = 0
+    missing_reward = 0
+    pred_missing = 0
+    seq_too_long = 0
+    sc_cnt = 0
+    # for item in tqdm(responses):
+    #     sorted_results = []
+    #     if not item["response"] or not item["pred"] or not item["res"]:
+    #         missing += 1
+    #         continue
+    #     for i, (resp, pred, r) in enumerate(zip(item["response"], item["pred"], item["res"])):
+    #         resp_id = f"{item['id']}_{i}"
+    #         if resp_id not in id2reward:
+    #             missing_reward += 1
+    #             continue
+    #
+    #         reward = id2reward[resp_id]
+    #         sorted_results.append((resp, pred, r, reward))
+    #
+    #     if not sorted_results:
+    #         continue
+    #     sorted_results = sorted(sorted_results, key=lambda x: x[-1], reverse=True)
+    #     for k in _k:
+    #         if any([r[2] for r in sorted_results[:k]]):
+    #             orm_pass_at_k[k] += 1
+    #
+    #     if item["sc_res"]:
+    #         sc_cnt += 1
+    with Pool(args.num_workers, initializer=_init, initargs=(id2reward,)) as pool:
+        annotate = partial(_worker, sc_top_k=(5, 10))
+        results = list(tqdm(pool.imap_unordered(annotate, responses), total=len(responses)))
+
+    sc_top_k = {k: 0 for k in (5, 10)}
+    for item in results:
+        missing += item["missing"]
+        missing_reward += item["reward_missing"]
+        pred_missing += item["pred_missing"]
+        seq_too_long += item["seq_too_long"]
+        sorted_results = item["sorted_results"]
+
+        if item["sc_res"]:
+            sc_cnt += 1
+
+        if not sorted_results:
+            continue
+
+        # ultimate_results.append(sorted_results)
+        for k in _k:
+            if any([r[2] for r in sorted_results[:k]]):
+                orm_pass_at_k[k] += 1
+
+        for k, v in item["sc_top_k_res"].items():
+            if v:
+                sc_top_k[k] += 1
+
+    print(f"Total: {len(responses)}")
+    print(f"Missing: {missing}")
+    print(f"Missing reward: {missing_reward}")
+    print(f"Missing pred: {pred_missing}")
+    print(f"Seq too long: {seq_too_long}")
+    print(f"SC: {sc_cnt}")
+    for k, v in orm_pass_at_k.items():
+        print(f"PRM pass at {k}: {v}")
+        print(f"PRM pass at {k} rate: {v / len(responses) * 100:.2f}%")
+    for k, v in sc_top_k.items():
+        print(f"SC pass at {k}: {v}")
+        print(f"SC pass at {k} rate: {v / len(responses) * 100:.2f}%")
+    print(f"SC rate: {sc_cnt / len(responses) * 100:.2f}%")
+
+
+if __name__ == '__main__':
+    main()