ask user whther to save pass@k

Author: marianna13

bigcodebench/evaluate.py

@@ -281,8 +281,28 @@ def stucking_checker():
     pass_at_k_path = result_path.replace("_eval_results.json", "_pass_at_k.json")
     pass_at_k["model"] = flags.samples.split("/")[-1].replace(".jsonl", "")
     pass_at_k["subset"] = flags.subset
-    with open(pass_at_k_path, "w") as f:
-        json.dump(pass_at_k, f, indent=2)
+
+    def save_pass_at_k():
+        with open(pass_at_k_path, "w") as f:
+            json.dump(pass_at_k, f, indent=2)
+
+    if os.path.isfile(pass_at_k_path):
+        saved_pass_at_k = json.load(open(pass_at_k_path, "r"))
+        # compare saved_pass_at_k with pass_at_k
+        for k in saved_pass_at_k.keys():
+            if pass_at_k[k] != saved_pass_at_k[k]:
+                cprint(f"Warning: {k} is different from the saved one", "yellow")
+                
+        # ask user whether to save the pass@k
+        decision = ""
+        while decision.lower() not in ["y", "n"]:
+            print(f"Save pass@k to {pass_at_k_path}? [Y/N]")
+            decision = input()
+        if decision.lower() == "y":
+            save_pass_at_k()
+            
+    else:
+        save_pass_at_k()
 
 
 def main():