add ignore assertion errors option

code_eval.py

@@ -168,6 +168,7 @@ class CodeEval(evaluate.Metric):
         timeout=3.0,
         early_stop=False,
         disable_tqdm=False,
+        ignore_assertion_errors=False,
     ):
         """Returns the scores"""
 
@@ -184,7 +185,11 @@ class CodeEval(evaluate.Metric):
             for tid, pred, ref in zip(ids, predictions, references):
                 results[tid] = []
                 for pid, p in enumerate(pred):
-                    result = Result(task_id=tid, prediction_id=pid)
+                    result = Result(
+                        task_id=tid,
+                        prediction_id=pid,
+                        ignore_assertion_error=ignore_assertion_errors,
+                    )
                     body = Template(pred_template).safe_substitute(prediction=p)
                     for r in ref:
                         assert isinstance(r, str)
@@ -192,7 +197,13 @@ class CodeEval(evaluate.Metric):
                         test = Template(test).safe_substitute(prediction=p)
 
                         test_program = body + "\n" + test
-                        args = (test_program, timeout, tid, pid)
+                        args = (
+                            test_program,
+                            timeout,
+                            tid,
+                            pid,
+                            ignore_assertion_errors,
+                        )
                         future = executor.submit(check_correctness, *args)
                         result.test_program.append(test_program)
                         result.add(future)
@@ -263,7 +274,8 @@ def estimate_pass_at_k(num_samples, num_correct, k):
 class Result(BaseModel):
     task_id: str
     prediction_id: int
-    
+    ignore_assertion_error: bool = False
+
     test_program: List[str] = []
     passed: Optional[bool] = None
     result: List[str] = []

execute.py

@@ -25,7 +25,13 @@ import signal
 import tempfile
 
 
-def check_correctness(check_program, timeout, task_id, completion_id):
+def check_correctness(
+    check_program,
+    timeout,
+    task_id,
+    completion_id,
+    ignore_assertion_errors=False,
+):
     """
     Evaluates the functional correctness of a completion by running the test
     suite provided in the problem.
@@ -36,7 +42,10 @@ def check_correctness(check_program, timeout, task_id, completion_id):
     manager = multiprocessing.Manager()
     result = manager.list()
 
-    p = multiprocessing.Process(target=unsafe_execute, args=(check_program, result, timeout))
+    p = multiprocessing.Process(
+        target=unsafe_execute,
+        args=(check_program, result, timeout, ignore_assertion_errors),
+    )
     p.start()
     p.join(timeout=timeout + 1)
     if p.is_alive():
@@ -53,10 +62,13 @@ def check_correctness(check_program, timeout, task_id, completion_id):
     )
 
 
-def unsafe_execute(check_program, result, timeout):
-
+def unsafe_execute(
+    check_program,
+    result,
+    timeout,
+    ignore_assertion_errors=False,
+):
     with create_tempdir():
-
         # These system calls are needed when cleaning up tempdir.
         import os
         import shutil
@@ -77,6 +89,11 @@ def unsafe_execute(check_program, result, timeout):
             result.append("passed")
         except TimeoutException:
             result.append("timed out")
+        except AssertionError as e:
+            if ignore_assertion_errors:
+                result.append("passed")
+            else:
+                result.append(f"failed: {e}")
         except BaseException as e:
             result.append(f"failed: {e}")
 
@@ -171,10 +188,16 @@ def reliability_guard(maximum_memory_bytes=None):
     if maximum_memory_bytes is not None:
         import resource
 
-        resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
-        resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
+        resource.setrlimit(
+            resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)
+        )
+        resource.setrlimit(
+            resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)
+        )
         if not platform.uname().system == "Darwin":
-            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+            resource.setrlimit(
+                resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)
+            )
 
     faulthandler.disable()