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,133 @@
+import copy
+import json
+import re
+import sys
+from argparse import ArgumentParser
+from datasets import load_dataset
+from collections import defaultdict
+import os
+
+from tqdm import tqdm
+
+sys.set_int_max_str_digits(0)
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+
+from scripts.apps.utils_execute import run_test
+
+
+def extract_content_between_tags(text):
+    # Regular expression to match content between [BEGIN] and [END]
+    pattern = r'\[BEGIN\](.*?)\[END\]'
+
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    else:
+        return None
+
+
+def extract_test_cases(text):
+    # Regular expression to match the contents between <TEST INPUT X> and </TEST INPUT X>
+    pattern = r'<TEST INPUT (\d+)>(.*?)</TEST INPUT \1>'
+
+    # Find all matches in the text, the re.DOTALL flag allows . to match newline characters
+    test_cases = re.findall(pattern, text, re.DOTALL)
+
+    # Extract only the content part from the matches
+    test_cases = [case[1].strip() for case in test_cases]
+
+    return test_cases
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument("--split", type=str, default="train")
+    parser.add_argument("--test_case_inputs", type=str, required=True, help="The completion file.")
+    parser.add_argument("--test_case_outputs", type=str, required=True, help="The completion file.")
+    # parser.add_argument("--output_file", type=str, required=True)
+    args = parser.parse_args()
+
+    data = load_dataset("codeparrot/apps", split="train").to_list()
+    print(len(data))
+    # for item in data:  # Currently, we do not require the solutions and input_output fields
+    #     if item["solutions"]:
+    #         item["solutions"] = json.loads(item["solutions"])
+    #     if item["input_output"]:
+    #         item["input_output"] = json.loads(item["input_output"])
+
+    case_inputs = [json.loads(line) for line in open(args.test_case_inputs).readlines()]
+    case_outputs = [json.loads(line) for line in open(args.test_case_outputs).readlines()]
+
+    problem_id2case = defaultdict(list)
+    for item in case_inputs:
+        cases = extract_test_cases(item["completion"])
+        if len(cases):
+            for i, case in enumerate(cases):
+                problem_id2case[item["problem_id"]].append({
+                    "input": case,
+                    "case_id": i,
+                })
+
+    for item in case_outputs:
+        output = extract_content_between_tags(item["completion"])
+        if output:
+            assert isinstance(output, str), output
+            case_id = int(item["case_id"].split("_")[1])
+            assert case_id == problem_id2case[item["problem_id"]][case_id]["case_id"]
+            problem_id2case[item["problem_id"]][case_id]["output"] = output
+
+    # Currently, we can directly verify
+    outputs = []
+    for item in data:
+        if item["solutions"]:
+            item["solutions"] = json.loads(item["solutions"])
+
+            problem_id = item["problem_id"]
+            cases = {
+                "inputs": [],
+                "outputs": [],
+            }
+            for case in problem_id2case[problem_id]:
+                if "output" in case:
+                    cases["inputs"].append(case["input"])
+                    cases["outputs"].append(case["output"])
+                    assert isinstance(case["input"], str), case["input"]
+                    assert isinstance(case["output"], str), case["output"]
+            if cases["inputs"]:
+                item["gen_test_cases"] = cases
+                outputs.append(item)
+
+    # print(outputs[0]["gen_test_cases"])
+
+    cnt = 0
+    num_program = 0
+    for i, item in tqdm(enumerate(outputs)):
+        if i in [7]:
+            continue  # Some bad cases.
+        print(item["gen_test_cases"])
+        for solution in item["solutions"]:
+            try:
+                res = run_test(item["gen_test_cases"], solution)
+                pass_num = sum([1 for item in res if item is True])
+                cnt += pass_num
+            except:
+                pass
+            num_program += 1
+
+    print(f"Pass rate: {cnt / num_program}")
+
+    # print(json.dumps(data[0], indent=2))
+    # print(data[0]["question"])
+    # print(data[0]["solutions"][0])
+    # print(data[0]["input_output"]["inputs"][0])
+    # print(data[0]["input_output"]["outputs"][0])
+
+    # with open(args.output_file, "w") as f:
+    #     for item in tqdm(outputs, desc="Writing prompts"):
+    #         f.write(json.dumps(item) + "\n")
+
+
+if __name__ == '__main__':
+    main()