fix(inspect): update inspect

Author: terryyz

bigcodebench/inspect.py

@@ -19,7 +19,7 @@ def inspection(args):
         shutil.rmtree(path, ignore_errors=True)
     if not os.path.exists(path):
         os.makedirs(path)
-    problems = get_bigcodebench()
+    problems = get_bigcodebench(subset=flags.subset)
 
     eval_results = json.load(open(args.eval_results, "r"))
     for task_id, results in eval_results["eval"].items():
@@ -30,7 +30,7 @@ def inspection(args):
             os.makedirs(task_path)
         task_id_data = problems[task_id]
         with open(os.path.join(task_path, "ground_truth.py"), "w") as f:
-            f.write(task_id_data[f"{args.subset}_prompt"] + "\n\n" + task_id_data["canonical_solution"])
+            f.write(task_id_data[f"{args.split}_prompt"] + "\n\n" + task_id_data["canonical_solution"])
 
         # write test
         with open(os.path.join(task_path, "test_case.py"), "w") as f:
@@ -49,7 +49,10 @@ def inspection(args):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--eval-results", required=True, type=str)
-    parser.add_argument("--subset", required=True, type=str)
+    parser.add_argument(
+        "--split", required=True, type=str, choices=["complete", "instruct"]
+    )
+    parser.add_argument("--subset", default="hard", type=str, choices=["full", "hard"])
     parser.add_argument("--in-place", action="store_true")
     args = parser.parse_args()