chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,154 @@
+import json
+import argparse
+from glob import glob
+import os
+import sys
+import collections
+from multiprocessing import Pool
+from tqdm import tqdm
+from functools import partial
+
+sys.set_int_max_str_digits(0)
+
+"""
+For self-consistency based input-output pairs, first copy the pseudo test cases into the prefix data, run `prefix_fail_extract_pseudo_label.py`,
+and the run this script.
+
+This script is incorrect for pseudo test cases since we cannot ensure the correctness of each test case, but it is appropriate for ground-truth test cases,
+serving as hard limit, i.e., if there is one completion for some prefix has passed all test cases, then it is a gold prefix.
+
+"""
+
+
+def counting_partial_response_value(res):
+    return sum([1 if x else 0 for x in res])
+
+
+def parse_value(v, binary: bool):
+    if binary:
+        return 1 if v > 0 else 0
+    return v
+
+
+def _process_trajectories_worker(item, top_k: int, binary: bool):
+    item_id, trajectories = item
+    outputs = {
+        "idx": item_id
+    }
+    for i in range(top_k):
+        level_trajectories = [(traj["vs"][i], traj["prefix"]) for traj in trajectories if len(traj["vs"]) > i]
+        if len(level_trajectories) == 0:
+            continue
+        prefix_vis = set()
+        level_values = []
+        level_prefixes = []
+        for v, prefix in level_trajectories:
+            if prefix in prefix_vis:
+                continue
+            level_values.append(parse_value(v, binary))
+            level_prefixes.append(prefix)
+            prefix_vis.add(prefix)
+        outputs[f"traj_level_{i}_values"] = level_values
+        outputs[f"traj_level_{i}_prefixes"] = level_prefixes
+    return outputs
+
+
+def _annotate(file):
+    return json.load(open(file, encoding="utf-8"))
+
+
+def multiprocessing_loading(files, num_workers: int = 8):
+    with Pool(num_workers) as p:
+        data = list(tqdm(p.imap(_annotate, files), total=len(files)))
+    all_data = []
+    for d in data:
+        all_data.extend(d)
+    return all_data
+    # data = []
+    # for file in tqdm(files):
+    #     data += json.load(open(file, encoding="utf-8"))
+    # return data
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_file", type=str)
+    parser.add_argument("--output_file", type=str)
+    parser.add_argument("--binary", default=False, action="store_true")
+    parser.add_argument("--num_workers", type=int, default=8)
+    args = parser.parse_args()
+
+    print("Collecting data...")
+    if os.path.exists(args.input_file):
+        data = json.load(open(args.input_file))
+    else:
+        # data = []
+        # for file in glob(args.input_file):
+        #     print(file)
+        #     data += json.load(open(file))
+        files = glob(args.input_file)
+        files = sorted(files)
+        print(len(files))
+        print(files)
+        data = multiprocessing_loading(files)
+
+    print(len(data))
+
+    num_prefixes = 0
+    val_cnt = collections.Counter()
+    outputs = []
+    preference_pairs = dict()
+    missing = 0
+    for item in tqdm(data):
+        problem_id, resp_id, prefix_id = item["prefix_id"].split("_")
+        prefix = item["prefix"]
+        problem_id = int(problem_id)
+
+        if "res" not in item:
+            missing += 1
+            continue
+        v = counting_partial_response_value(item["res"])
+        v = parse_value(v, args.binary)
+
+        outputs.append({
+            "problem_id": problem_id,
+            "prefix": prefix,
+            "value": v,
+        })
+        num_prefixes += 1
+        val_cnt[v] += 1
+
+        if problem_id not in preference_pairs:
+            preference_pairs[problem_id] = {
+                "pos": [],
+                "neg": [],
+            }
+
+        if v > 0:
+            preference_pairs[problem_id]["pos"].append(prefix)
+        else:
+            preference_pairs[problem_id]["neg"].append(prefix)
+
+    preference_pairs = [
+        {
+            "problem_id": problem_id,
+            "pos": pair["pos"],
+            "neg": pair["neg"],
+        }
+        for problem_id, pair in preference_pairs.items()
+    ]
+
+    print(f"Missing: {missing}")
+    print(val_cnt)
+    print(f"Processed {num_prefixes} prefixes.")
+    print(f"Averaged {num_prefixes / len(data)} prefixes per problem.")
+    json.dump(outputs, open(args.output_file, "w", encoding="utf-8"), indent=2, ensure_ascii=False)
+    json.dump(preference_pairs, open(args.output_file.replace(".json", "_pairs.json"), "w", encoding="utf-8"), indent=2, ensure_ascii=False)
+
+
+if __name__ == '__main__':
+    main()
+
+"""
+>>> 
+"""