timed re.search and _executor made global

Author: francesco-bertolotti

opencompass/datasets/mbpp.py

@@ -4,7 +4,7 @@
 import json
 import multiprocessing
 import os.path as osp
-import re
+import regex as re
 import signal
 import tempfile
 from collections import defaultdict
@@ -330,7 +330,7 @@ def _process_answer(self, text):
             r"'(.*)'\s*\[DONE\]",
         ]
         for p in patterns:
-            match = re.search(p, text, re.DOTALL)
+            match = re.search(p, text, re.DOTALL, timeout=10.0)
             if match:
                 text = match.group(1)
                 break
@@ -383,6 +383,22 @@ def _process_answer(self, text):
             text = text[1:]
         return text
 
+def _execution(programs, timeout, key):
+    try:
+        # Add exec globals to prevent the exec to raise
+        # unnecessary NameError for correct answer
+        exec_globals = {}
+        with swallow_io():
+            with time_limit(timeout):
+                exec(programs, exec_globals)
+        key.append('pass')
+    except TimeOutException:
+        key.append('timeout')
+    except AssertionError:
+        key.append('wrong_answer')
+    except BaseException as e:
+        print(e)
+        key.append('failed')
 
 def execution(programs, task_id, timeout):
     """Execution function for running generation code.
@@ -400,29 +416,12 @@ def execution(programs, task_id, timeout):
     control the process.
     """
 
-    def _execution(programs, timeout):
-        try:
-            # Add exec globals to prevent the exec to raise
-            # unnecessary NameError for correct answer
-            exec_globals = {}
-            with swallow_io():
-                with time_limit(timeout):
-                    exec(programs, exec_globals)
-            key.append('pass')
-        except TimeOutException:
-            key.append('timeout')
-        except AssertionError:
-            key.append('wrong_answer')
-        except BaseException as e:
-            print(e)
-            key.append('failed')
-
     manager = multiprocessing.Manager()
     key = manager.list()
     # `signal` cannot be used in child thread, therefore, we
     # need to create a process in the thread.
     p = multiprocessing.Process(target=_execution,
-                                args=(programs, timeout - 1))
+                                args=(programs, timeout - 1, key))
     p.start()
     p.join(timeout=timeout)
     if p.is_alive():