Update code_efficiency_calculator.py

code_efficiency_calculator.py

@@ -19,7 +19,32 @@ import io
 import os
 import signal
 from tqdm import tqdm
+import ast
 
+def is_string_expression(node):
+    """检查节点是否为字符串表达式（Str或Constant节点且值为字符串）"""
+    if isinstance(node, ast.Str):
+        return True
+    elif isinstance(node, ast.Constant):
+        return isinstance(node.value, str)
+    return False
+
+def is_only_imports_and_function_headers(code):
+    try:
+        tree = ast.parse(code)
+    except SyntaxError:
+        return False
+    
+    for node in tree.body:
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            continue
+        elif isinstance(node, ast.FunctionDef):
+            for stmt in node.body:
+                if not (isinstance(stmt, ast.Expr) and is_string_expression(stmt.value)):
+                    return False
+        else:
+            return False
+    return True
 
 class TimeoutException(Exception):
     pass
@@ -388,7 +413,7 @@ def calculate_line_efficiency(completion_file,entry_point):
             # print(report_content)
 
     except subprocess.CalledProcessError as e:
-        # print(f"Error during the execution: {e}")
+        print(f"Error during the execution: {e}")
         report_content = f"Error during the execution: {e}"
 
     # # 清理临时文件
@@ -417,14 +442,36 @@ def humaneval_add_string_to_py_file(data,evaluation_code=False, path="./tmp/"):
         else:
             code.append(string)
     data["completion"] = "\n".join(code)
+    solution1 = data["completion"]
+    flagg = 0
     try:
-        if f"```python" in data["completion"]:
-            start_idx = data["completion"].find(f"```python")
-            data["completion"] = data["completion"][start_idx+len(f"```python"):]
-            if "```" in data["completion"]:
-                end_idx = data["completion"].find("```")
-                data["completion"] = data["completion"][:end_idx]
-        full_code = import_pkg+ "\n"+data["prompt"] + "\n"+data["completion"] + "\n" + test_case
+        try: 
+            while("```python" in solution1):              
+                start_idx = solution1.find(f"```python")
+                solution1 = solution1[start_idx+len(f"```python"):]
+                if "```" in solution1:
+                    end_idx = solution1.find("```")
+                    if  data["entry_point"] in solution1[:end_idx]:
+                        if is_only_imports_and_function_headers(solution1[:end_idx]):
+                            solution1 = solution1[end_idx+len(f"```"):]
+                        else:
+                            data["completion"] = solution1[:end_idx]
+                            flagg = 1
+                            break
+                    else:
+                        solution1 = solution1[end_idx+len(f"```"):]
+                else:
+                    break
+        except:
+          pass
+        if flagg == 0:
+            if f"```python" in data["completion"]:
+                start_idx = data["completion"].find(f"```python")
+                data["completion"] = data["completion"][start_idx+len(f"```python"):]
+                if "```" in data["completion"]:
+                    end_idx = data["completion"].find("```")
+                    data["completion"] = data["completion"][:end_idx]
+        full_code = import_pkg+ "\n"+data["prompt"] + "\n"+data["completion"] + "\n" + test_case +"\n"+f"check({data['entry_point']})"
         # with open(f"./{path}/{problem_idx}.py", "w") as f:
         #     f.write(full_code)
         # return_path = f"./{path}/{problem_idx}.py"
@@ -477,7 +524,7 @@ def mbpp_add_string_to_py_file(data,evaluation_code=False, path="./tmp/"):
                 f.write(full_code)
             return_path = f"./{path}/{problem_idx}.py"
     except Exception as e:
-        # print(e)
+        print(e)
         pass
     # print(return_path,full_code)
     return return_path,full_code
@@ -524,7 +571,7 @@ def add_string_to_py_file(data,evaluation_code=False, path="./tmp/"):
             else:
                 return_path = None
     except Exception as e:
-        # print(e)
+        print(e)
         pass
     return return_path,full_code
 
@@ -554,7 +601,7 @@ def calculate_code_execution_efficiency(data,evaluation_code=False,path="./tmp/"
             problem_idx = data["problem_idx"]
             completion_file,full_code = add_string_to_py_file(data,evaluation_code=evaluation_code, path=path)
     except Exception as e:
-        # print(e)
+        print(e)
         completion_file = None
     if completion_file == None:
         # print("test")
@@ -583,7 +630,7 @@ The total execution time is: {canonical_solution_execution_time} s.
 The maximum memory peak requirement is: {canonical_solution_max_memory_usage} MB.
 """
     except Exception as e:
-        # print(e)
+        print(e)
         overhead = f"""
 The code execution failed.
 """