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,68 @@
+import os.path
+
+from datasets import load_dataset
+import argparse
+import json
+from tqdm import tqdm
+from glob import glob
+
+import re
+
+
+def extract_test_case_inputs(text):
+    # Regular expression to match content between <TEST CASE INPUTS> and </TEST CASE INPUTS>
+    pattern = r'<TEST CASE INPUTS>(.*?)</TEST CASE INPUTS>'
+
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    else:
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_file", type=str)
+    parser.add_argument("--output_file", type=str)
+    args = parser.parse_args()
+
+    data = []
+    if os.path.exists(args.input_file):
+        data.extend([json.loads(line) for line in open(args.input_file).readlines()])
+    else:
+        for file in glob(args.input_file):
+            data.extend([json.loads(line) for line in open(file).readlines()])
+
+    outputs = {}
+    for item in data:
+        cases = extract_test_case_inputs(item["completion"])
+        if cases is not None:
+            cases = cases.split("\n")
+            cases = [case.strip() for case in cases if case.strip()]
+            cases = [case for case in cases if "==" not in case]
+            cases = [case for case in cases if "assert " in case]
+
+            if item["task_id"] not in outputs:
+                outputs[item["task_id"]] = {
+                    "source_file": item["source_file"],
+                    "task_id": item["task_id"],
+                    "code": item["code"],
+                    "test_imports": item["test_imports"],
+                    "test_list": item["test_list"],
+                }
+                outputs[item["task_id"]]["aug_cases"] = set()
+
+            outputs[item["task_id"]]["aug_cases"].update(cases)
+
+    cnt = 0
+    for item_id in outputs:
+        outputs[item_id]["aug_cases"] = list(outputs[item_id]["aug_cases"])
+        cnt += len(outputs[item_id]["aug_cases"])
+
+    print(f"Averaged {cnt/len(outputs)} test cases per task.")
+
+    json.dump(outputs, open(args.output_file, "w"), indent=2)
+
+
+if __name__ == "__main__":
+    main()