update

calculate_memory_usage.py

@@ -0,0 +1,220 @@
+import json
+import os
+import glob
+import numpy as np
+import argparse
+from code_efficiency_calculator import run_model_task
+
+
+def calculate_memory_usage(dat_file_path):
+    with open(dat_file_path, 'r') as file:
+        prev_time = 0
+        prev_mem_mb = 0
+        mem_time_mb_s = 0
+        next(file)
+        for line in file:
+            if "__main__." in line:
+                continue
+            parts = line.split()
+            mem_in_mb = float(parts[1])
+            timestamp = float(parts[2])
+            if prev_time > 0:
+                time_interval_s = timestamp - prev_time
+                mem_time_mb_s += (prev_mem_mb + mem_in_mb) / 2 * time_interval_s
+            prev_time = timestamp
+            prev_mem_mb = mem_in_mb
+        return mem_time_mb_s
+
+
+def calculate_runtime(dat_file_path):
+    with open(dat_file_path, 'r') as file:
+        start_time = float("inf")
+        end_time = float("-inf")
+        next(file)
+        for line in file:
+            if "__main__." in line:
+                continue
+            parts = line.split()
+            timestamp = float(parts[2])
+            start_time = min(start_time, timestamp)
+            end_time = max(end_time, timestamp)
+        return max(end_time - start_time,0)
+
+def report_max_memory_usage(dat_file_path):
+    max_memory_usage = 0
+    with open(dat_file_path, 'r') as file:
+        next(file)
+        for line in file:
+            if "__main__." in line:
+                continue
+            parts = line.split()
+            mem_in_mb = float(parts[1])
+            max_memory_usage = max(max_memory_usage, mem_in_mb)
+        return max_memory_usage
+
+def report_results(task,model):
+    run_model_task(task,model)
+    dat_directory = f"./results/{task}_{model}"
+    canonical_solution_directory = f"./results/{task}_canonical_solution"
+    canonical_solution_memory_usage = {}
+    canonical_solution_execution_time = {}
+    canonical_solution_max_memory_usage = {}
+    for dat_file in glob.glob(os.path.join(canonical_solution_directory, "*.dat")):
+        try:
+            problem_idx = os.path.basename(dat_file).split('.')[0]
+            canonical_solution_memory_usage[int(problem_idx)] = calculate_memory_usage(dat_file)
+            canonical_solution_execution_time[int(problem_idx)] = calculate_runtime(dat_file)
+            canonical_solution_max_memory_usage[int(problem_idx)] = report_max_memory_usage(dat_file)
+        except:
+            pass
+
+
+    global_result = {}
+
+
+    completion_memory_usage = {}
+    execution_time = {}
+    max_memory_usage = {}
+    task_idx = {}
+    for dat_file in glob.glob(os.path.join(dat_directory, "*.dat")):
+        try:
+            problem_idx = os.path.basename(dat_file).split('.')[0]
+            execution_time_result = calculate_runtime(dat_file)
+            completion_memory_usage[int(problem_idx)] = calculate_memory_usage(dat_file)
+            execution_time[int(problem_idx)] = calculate_runtime(dat_file)
+            max_memory_usage[int(problem_idx)] = report_max_memory_usage(dat_file)
+            task_idx[int(problem_idx)] = dat_file
+        except Exception as e:
+            print(dat_file)
+    global_result[model] = {"completion_memory_usage":completion_memory_usage,"execution_time":execution_time,"max_memory_usage":max_memory_usage,"task_idx":task_idx}
+
+
+
+    save_results = []
+    max_net_lists = {}
+    max_nmu_lists = {}
+    max_ntmu_lists = {}
+
+    for model in global_result.keys():
+        completion_memory_usage = global_result[model]["completion_memory_usage"]
+        execution_time = global_result[model]["execution_time"]
+        max_memory_usage = global_result[model]["max_memory_usage"]
+
+        # report execution time
+        total_execution_time = 0
+
+        # report normalized execution time
+        normalized_execution_time = 0
+
+        # report max memory usage
+        total_max_memory_usage = 0
+
+        # report normalized max memory usage
+        normalized_max_memory_usage = 0
+
+        # report memory usage
+        total_memory_usage = 0
+        total_canonical_solution_max_memory_usage = 0
+        total_canonical_solution_execution_time = 0
+        total_canonical_solution_memory_usage = 0
+        # report normalized memory usage
+        normalized_memory_usage = 0
+        total_codes = 0
+        normalized_execution_time_list = []
+        normalized_max_memory_usage_list = []
+        normalized_memory_usage_list = []
+        total_fast = 0
+        total_95 = 0
+        total_97=0
+        total_99=0
+        total_100=0
+        total_101=0
+        total_1000=0
+        total_500=0
+        category_tmp = {}
+        total_10000=0
+        max_net = float("-inf")
+        max_nmu = float("-inf")
+        max_tmu = float("-inf")
+
+        total_500_net = 0
+        total_500_nmu = 0
+        total_500_tmu = 0
+        # print(len(completion_memory_usage))
+        for idx in completion_memory_usage.keys():
+            if idx not in canonical_solution_memory_usage.keys():
+                continue
+
+
+            total_memory_usage += completion_memory_usage[idx]
+            total_execution_time += execution_time[idx]
+            total_max_memory_usage += max_memory_usage[idx]
+            total_canonical_solution_max_memory_usage+=canonical_solution_max_memory_usage[idx]
+            total_canonical_solution_memory_usage+=canonical_solution_memory_usage[idx]
+            total_canonical_solution_execution_time+=canonical_solution_execution_time[idx]
+            if execution_time[idx]/canonical_solution_execution_time[idx]>5:
+                total_500_net+=1
+            if max_net<execution_time[idx]/canonical_solution_execution_time[idx]:
+                max_net = execution_time[idx]/canonical_solution_execution_time[idx]
+            normalized_execution_time += execution_time[idx]/canonical_solution_execution_time[idx]
+            normalized_execution_time_list.append(execution_time[idx]/canonical_solution_execution_time[idx])
+            if max_memory_usage[idx]/canonical_solution_max_memory_usage[idx]>5:
+                total_500_nmu+=1
+            if max_nmu<max_memory_usage[idx]/canonical_solution_max_memory_usage[idx]:
+                max_nmu = max_memory_usage[idx]/canonical_solution_max_memory_usage[idx]
+            normalized_max_memory_usage += max_memory_usage[idx]/canonical_solution_max_memory_usage[idx]
+            normalized_max_memory_usage_list.append(max_memory_usage[idx]/canonical_solution_max_memory_usage[idx])
+
+            if completion_memory_usage[idx]/canonical_solution_memory_usage[idx]>5:
+                total_500_tmu+=1
+            net = execution_time[idx] / canonical_solution_execution_time[idx]
+            nmu = completion_memory_usage[idx] / canonical_solution_memory_usage[idx]
+            ntmu = max_memory_usage[idx] / canonical_solution_max_memory_usage[idx]
+            normalized_memory_usage += completion_memory_usage[idx]/canonical_solution_memory_usage[idx]
+            normalized_memory_usage_list.append(completion_memory_usage[idx]/canonical_solution_memory_usage[idx])
+
+            if len(max_net_lists) < 10 or net > min(max_net_lists.keys()):
+                if len(max_net_lists) >= 10:
+                    min_key = min(max_net_lists.keys())
+                    del max_net_lists[min_key]
+                max_net_lists[net] = (model, idx)
+
+            if len(max_nmu_lists) < 10 or nmu > min(max_nmu_lists.keys()):
+                if len(max_nmu_lists) >= 10:
+                    min_key = min(max_nmu_lists.keys())
+                    del max_nmu_lists[min_key]
+                max_nmu_lists[nmu] = (model, idx)
+
+            if len(max_ntmu_lists) < 10 or ntmu > min(max_ntmu_lists.keys()):
+                if len(max_ntmu_lists) >= 10:
+                    min_key = min(max_ntmu_lists.keys())
+                    del max_ntmu_lists[min_key]
+                max_ntmu_lists[ntmu] = (model, idx)
+            max_tmu = max(max_tmu,completion_memory_usage[idx]/canonical_solution_memory_usage[idx])
+            total_codes+=1
+
+        if len(normalized_execution_time_list)==0:
+            print(model)
+            continue
+        normalized_execution_time = normalized_execution_time/len(normalized_execution_time_list)
+        normalized_max_memory_usage = normalized_max_memory_usage/len(normalized_execution_time_list)
+        normalized_memory_usage = normalized_memory_usage/len(normalized_execution_time_list)
+        total_execution_time = total_execution_time/len(normalized_execution_time_list)
+        total_memory_usage = total_memory_usage/len(normalized_execution_time_list)
+        total_max_memory_usage = total_max_memory_usage/len(normalized_execution_time_list)
+
+        pass1 = len(completion_memory_usage)/1000*100
+
+        total_500_net = total_500_net/len(normalized_execution_time_list)*100
+        total_500_nmu = total_500_nmu/len(normalized_execution_time_list)*100
+        total_500_tmu = total_500_tmu/len(normalized_execution_time_list)*100
+
+        print(f"{model}&{total_execution_time:.2f}&{normalized_execution_time:.2f}&{max_net:.2f}&{total_500_net:.1f}&{total_max_memory_usage:.2f}&{normalized_max_memory_usage:.2f}&{max_nmu:.2f}&{total_500_nmu:.1f}&{total_memory_usage:.2f}&{normalized_memory_usage:.2f}&{max_tmu:.2f}&{total_500_tmu:.1f}&{pass1:.1f}\\\\")
+
+if __name__ == "__main__":
+    parse = argparse.ArgumentParser()
+    parse.add_argument("--task", type=str, default="EffiBench")
+    parse.add_argument("--model", type=str, default="gpt-4")
+
+    args = parse.parse_args()
+    report_results(args.task,args.model)

code_efficiency_calculator.py

@@ -0,0 +1,655 @@
+import json
+import os
+from execution import check_correctness
+import copy
+import argparse
+from tqdm import tqdm
+import subprocess
+from concurrent.futures import ThreadPoolExecutor
+import concurrent.futures
+import os
+import re
+import shutil
+import contextlib
+import random
+import json
+from typing import Optional, Callable, Dict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import inspect
+import numpy as np
+import sys
+import concurrent.futures
+import time
+from tqdm import tqdm
+import contextlib
+import faulthandler
+import io
+import os
+import multiprocessing
+import platform
+import signal
+from tqdm import tqdm
+
+
+class TimeoutException(Exception):
+    pass
+class WriteOnlyStringIO(io.StringIO):
+    """ StringIO that throws an exception when it's read from """
+
+    def read(self, *args, **kwargs):
+        raise IOError
+
+    def readline(self, *args, **kwargs):
+        raise IOError
+
+    def readlines(self, *args, **kwargs):
+        raise IOError
+
+    def readable(self, *args, **kwargs):
+        """ Returns True if the IO object can be read. """
+        return False
+class redirect_stdin(contextlib._RedirectStream):  # type: ignore
+    _stream = 'stdin'
+
+@contextlib.contextmanager
+def swallow_io():
+    stream = WriteOnlyStringIO()
+    with contextlib.redirect_stdout(stream):
+        with contextlib.redirect_stderr(stream):
+            with redirect_stdin(stream):
+                yield
+
+@contextlib.contextmanager
+def time_limit(seconds: float):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+
+ListNode_text = """
+class ListNode:
+    def __init__(self, val=0, next=None):
+        self.val = val
+        self.next = next
+"""
+TreeNode_text = """
+class TreeNode:
+    def __init__(self, val=0, left=None, right=None, next=None):
+        self.val = val
+        self.left = left
+        self.right = right
+        self.next = next
+"""
+
+import_pkg = """
+from typing import *
+from bisect import *
+from collections import *
+from copy import *
+from datetime import *
+from heapq import *
+from math import *
+from re import *
+from string import *
+from random import *
+from itertools import *
+from functools import *
+from operator import *
+
+import string
+import re
+import datetime
+import collections
+import heapq
+import bisect
+import copy
+import math
+import random
+import itertools
+import functools
+import operator
+"""
+
+memory_profiler_prompt = r"""
+def parse_profile_table(profile_table: str):
+    table = {"filename": None, "rows": []}
+    for line in profile_table.strip().split("\n"):
+        if line.startswith("Filename:"):
+            table["filename"] = line.split(": ")[1]
+        elif re.match(r"^\s*\d+", line):
+            parts = re.split(r"\s{2,}", line.strip(), maxsplit=4)
+            if len(parts) == 5 and "iB" in parts[1] and "iB" in parts[2]:
+                table["rows"].append({
+                    "line": int(parts[0]),
+                    "mem_usage": parts[1],
+                    "increment": parts[2],
+                    "occurrences": int(parts[3]),
+                    "line_contents": parts[4],
+                })
+            else:
+                parts = re.split(r"\s{2,}", line.strip(), maxsplit=1)
+                table["rows"].append({
+                    "line": int(parts[0]),
+                    "line_contents": parts[1] if len(parts) == 2 else "",
+                })
+    return table
+
+def print_averaged_results(profile_log: str, precision: int = 1):
+    tables = [parse_profile_table(table) for table in profile_log.split("\n\n\n")]
+    averaged_table = defaultdict(lambda: defaultdict(list))
+
+    for table in tables:
+        filename = table["filename"]
+        for row in table["rows"]:
+            line = row["line"]
+            if "mem_usage" in row:
+                mem_usage = float(row["mem_usage"].split()[0])
+                increment = float(row["increment"].split()[0])
+                occurrences = row["occurrences"]
+                averaged_table[filename][line].append((mem_usage, increment, occurrences))
+            else:
+                averaged_table[filename][line].append(tuple())
+
+    stream = sys.stdout
+    template = '{0:>6} {1:>12} {2:>12}  {3:>10}   {4:<}'
+
+    for filename, lines in averaged_table.items():
+        header = template.format('Line #', 'Mem usage', 'Increment', 'Occurrences', 'Line Contents')
+
+        stream.write(u'Filename: ' + filename + '\n\n')
+        stream.write(header + u'\n')
+        stream.write(u'=' * len(header) + '\n')
+
+        all_lines = linecache.getlines(filename)
+
+        float_format = u'{0}.{1}f'.format(precision + 4, precision)
+        template_mem = u'{0:' + float_format + '} MiB'
+
+        for lineno, mem_values in lines.items():
+            # TODO: should average the rest or not?
+            # mem_values = [(50.1, 0.0, 4), (51.1, 0.0, 6), ()]
+            if any([len(m) == 0 for m in mem_values]):
+                tmp = template.format(lineno, "", "", "", all_lines[lineno - 1])
+            else:
+                mem_usage_sum = sum(m[0] for m in mem_values)
+                increment_sum = sum(m[1] for m in mem_values)
+                occurrences_sum = sum(m[2] for m in mem_values)
+                count = len(mem_values)
+
+                avg_mem_usage = mem_usage_sum / count
+                avg_increment = increment_sum / count
+                avg_occurrences = occurrences_sum / count
+
+                avg_mem_usage_str = template_mem.format(avg_mem_usage)
+                avg_increment_str = template_mem.format(avg_increment)
+
+                tmp = template.format(lineno, avg_mem_usage_str, avg_increment_str, int(avg_occurrences), all_lines[lineno - 1])
+            stream.write(tmp)
+
+print_averaged_results(profile_stream.getvalue(), precision=PROFILE_PRECISION)
+"""
+
+memory_profiler_pkgs = r"""
+from collections import defaultdict, deque
+from memory_profiler import profile
+import io
+profile_stream = io.StringIO()
+PROFILE_PRECISION = 1
+import re
+import sys
+import linecache
+"""
+
+
+def calculate_memory_usage(dat_file_path):
+    with open(dat_file_path, 'r') as file:
+        prev_time = 0
+        prev_mem_mb = 0
+        mem_time_mb_s = 0
+        next(file)
+        for line in file:
+            if not line.startswith('MEM'):
+                continue  # Skip any line that does not start with 'MEM'
+            parts = line.split()
+            mem_in_mb = float(parts[1])
+            timestamp = float(parts[2])
+            if prev_time > 0:
+                time_interval_s = timestamp - prev_time
+                mem_time_mb_s += (prev_mem_mb + mem_in_mb) / 2 * time_interval_s
+            prev_time = timestamp
+            prev_mem_mb = mem_in_mb
+        return mem_time_mb_s
+
+
+def calculate_runtime(dat_file_path):
+    with open(dat_file_path, 'r') as file:
+        start_time = float("inf")
+        end_time = float("-inf")
+        next(file)
+        for line in file:
+            if not line.startswith('MEM'):
+                continue  # Skip any line that does not start with 'MEM'
+            parts = line.split()
+            timestamp = float(parts[2])
+            start_time = min(start_time, timestamp)
+            end_time = max(end_time, timestamp)
+        return max(end_time - start_time,0)
+
+def report_max_memory_usage(dat_file_path):
+    max_memory_usage = 0
+    with open(dat_file_path, 'r') as file:
+        prev_time = 0
+        prev_mem_mb = 0
+        mem_time_mb_s = 0
+        next(file)
+        for line in file:
+            if not line.startswith('MEM'):
+                continue  # Skip any line that does not start with 'MEM'
+            parts = line.split()
+            mem_in_mb = float(parts[1])
+            max_memory_usage = max(max_memory_usage, mem_in_mb)
+        return max_memory_usage
+
+def add_profile_decorator_to_python_file(file_path,entry_point):
+    """给Python文件中的函数自动添加@profile装饰器。"""
+    try:
+        with open(file_path, 'r') as file:
+            lines = file.readlines()
+        if "humaneval" in file_path:
+            with open(file_path, 'w') as file:
+                inside_class = False
+                class_indent = 0
+                for line in lines:
+                    stripped_line = line.lstrip()
+                    if stripped_line.startswith(f"def {entry_point}"):
+                        inside_class = True
+                        class_indent = len(line) - len(stripped_line)
+                        file.write('@profile\n')
+                        file.write(line)
+                        continue
+                    if inside_class:
+                        if stripped_line and not line[class_indent].isspace():
+                            inside_class = False
+                        elif stripped_line.startswith("def "):
+                            file.write(' ' * class_indent + '@profile\n')
+                    file.write(line)
+        if "mbpp" in file_path:
+            entry_point
+            with open(file_path, 'w') as file:
+                inside_class = False
+                class_indent = 0
+                for line in lines:
+                    stripped_line = line.lstrip()
+                    if stripped_line.startswith(f"def {entry_point}"):
+                        inside_class = True
+                        class_indent = len(line) - len(stripped_line)
+                        file.write('@profile\n')
+                        file.write(line)
+                        continue
+                    if inside_class:
+                        if stripped_line and not line[class_indent].isspace():
+                            inside_class = False
+                        elif stripped_line.startswith("def "):
+                            file.write(' ' * class_indent + '@profile\n')
+                    file.write(line)
+        else:
+            with open(file_path, 'w') as file:
+                inside_class = False
+                class_indent = 0
+                for line in lines:
+                    stripped_line = line.lstrip()
+                    if stripped_line.startswith("class Solution"):
+                        inside_class = True
+                        class_indent = len(line) - len(stripped_line)
+                        file.write(line)
+                        continue
+                    if inside_class:
+                        if stripped_line and not line[class_indent].isspace():
+                            inside_class = False
+                        elif stripped_line.startswith("def "):
+                            file.write(' ' * class_indent + '    @profile\n')
+                    file.write(line)
+    except Exception as e:
+        # print(f"Error during the file processing: {e}")
+        pass
+
+def add_profile_for_memory_profiler(code_string,data):
+    """给Python代码中的函数自动添加@profile装饰器。"""
+    entry_point = ""
+    try:
+        if "task_id" in data.keys() and "HumanEval" in data["task_id"]:
+            entry_point = data["entry_point"]
+            lines = code_string.split('\n')
+            new_lines = []
+            inside_class = False
+            class_indent = 0
+            first_function = True
+            for line in lines:
+                stripped_line = line.lstrip()
+                if stripped_line.startswith(f"def {entry_point}"):
+                    inside_class = True
+                    class_indent = len(line) - len(stripped_line)
+                    new_lines.append(' ' * class_indent + '@profile(stream=profile_stream, precision=PROFILE_PRECISION)')
+                new_lines.append(line)
+            return '\n'.join(new_lines)
+        elif "task_id" in data.keys():
+            entry_point = data["entry_point"]
+            lines = code_string.split('\n')
+            new_lines = []
+            inside_class = False
+            class_indent = 0
+            first_function = True
+            for line in lines:
+                stripped_line = line.lstrip()
+                if stripped_line.startswith(f"def {entry_point}"):
+                    inside_class = True
+                    class_indent = len(line) - len(stripped_line)
+                    new_lines.append(' ' * class_indent + '@profile(stream=profile_stream, precision=PROFILE_PRECISION)')
+                new_lines.append(line)
+            return '\n'.join(new_lines)
+        else:
+            lines = code_string.split('\n')
+            new_lines = []
+            inside_class = False
+            class_indent = 0
+            first_function = True
+            for line in lines:
+                stripped_line = line.lstrip()
+                if stripped_line.startswith("class Solution"):
+                    inside_class = True
+                    class_indent = len(line) - len(stripped_line)
+                    new_lines.append(line)
+                    continue
+                if inside_class:
+                    if stripped_line and not line[class_indent].isspace():
+                        inside_class = False
+                    elif stripped_line.startswith("def ") and first_function:
+                        new_lines.append(' ' * class_indent + '    @profile(stream=profile_stream, precision=PROFILE_PRECISION)')
+                        first_function = False
+                new_lines.append(line)
+            return '\n'.join(new_lines)
+    except Exception as e:
+        return code_string
+
+def calculate_line_efficiency(completion_file,entry_point):
+    try:
+        path, filename = os.path.split(completion_file)
+        tmp_py_script_filename = f"{filename.split('.')[0]}_tmp.py"
+        tmp_py_script = os.path.join(path, tmp_py_script_filename)
+        tmp_lprof_filename = f"{tmp_py_script_filename}.lprof"  # 期望的lprof文件名
+        
+        # 复制原始脚本到临时文件，并添加@profile装饰器
+        subprocess.run(['cp', completion_file, tmp_py_script],check=True, capture_output=True, text=True)
+        add_profile_decorator_to_python_file(tmp_py_script,entry_point)
+
+        subprocess.run(['timeout',"10",'kernprof', '-l', tmp_py_script_filename], cwd=path, capture_output=True, text=True, check=True)
+        # 生成性能报告
+        overhead_dir = path
+        # os.makedirs(overhead_dir, exist_ok=True)
+        report_file = os.path.join(overhead_dir, tmp_py_script_filename.replace('.py', '.txt'))
+        with open(report_file, 'w') as f:
+            subprocess.run(['timeout',"10",'python', '-m', 'line_profiler', tmp_lprof_filename], cwd=path, stdout=f)
+        with open(report_file, 'r') as f:
+            report_content = f.read()
+            # print(report_content)
+
+    except subprocess.CalledProcessError as e:
+        # print(f"Error during the execution: {e}")
+        report_content = f"Error during the execution: {e}"
+
+    # # 清理临时文件
+    if os.path.exists(tmp_py_script):
+        os.remove(tmp_py_script)
+    if os.path.exists(f"{tmp_py_script}.lprof"):
+        os.remove(f"{tmp_py_script}.lprof")
+
+    return report_content
+
+def humaneval_add_string_to_py_file(data,evaluation_code=False, path="./tmp/"):
+    if "canonical_solution" in path:
+        data["completion"] = data["canonical_solution"]
+    if evaluation_code==False:
+        test_case = data["test"]
+    else: 
+        test_case = data["small_test_cases"]
+    # test_case = data["small_test_cases"]
+    problem_idx = data["task_id"].split("/")[1]
+    return_path,full_code = None,""
+    tmp_code = data["completion"].split("\n")
+    code = []
+    for string in tmp_code:
+        if "print(" in string:
+            continue
+        else:
+            code.append(string)
+    data["completion"] = "\n".join(code)
+    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
+        # with open(f"./{path}/{problem_idx}.py", "w") as f:
+        #     f.write(full_code)
+        # return_path = f"./{path}/{problem_idx}.py"
+        result = check_correctness(full_code,timeout=10.0)
+        if result["passed"]:
+            with open(f"./{path}/{problem_idx}.py", "w") as f:
+                f.write(full_code)
+            return_path = f"./{path}/{problem_idx}.py"
+            # print(return_path)
+        else:
+            return_path = None
+    except Exception as e:
+        pass
+    # print(return_path,full_code)
+    return return_path,full_code
+
+
+def mbpp_add_string_to_py_file(data,evaluation_code=False, path="./tmp/"):
+    if "canonical_solution" in path:
+        data["completion"] = data["code"]
+    if evaluation_code==False:
+        test_case = data["test"]
+    else: 
+        test_case = "\n".join(data["test_list"])
+    # test_case = data["small_test_cases"]
+    problem_idx = str(data["task_id"])
+    return_path,full_code = None,""
+    tmp_code = data["completion"].split("\n")
+    code = []
+    for string in tmp_code:
+        if "print(" in string:
+            continue
+        else:
+            code.append(string)
+    data["completion"] = "\n".join(code)
+    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 = "\n".join(data["test_imports"])+ "\n"+data["completion"] + "\n" + test_case
+        # with open(f"./{path}/{problem_idx}.py", "w") as f:
+        #     f.write(full_code)
+        # return_path = f"./{path}/{problem_idx}.py"
+        result = check_correctness(full_code,timeout=10.0)
+        if result["passed"]:
+            with open(f"./{path}/{problem_idx}.py", "w") as f:
+                f.write(full_code)
+            return_path = f"./{path}/{problem_idx}.py"
+    except Exception as e:
+        # print(e)
+        pass
+    # print(return_path,full_code)
+    return return_path,full_code
+
+def add_string_to_py_file(data,evaluation_code=False, path="./tmp/"):
+    if "canonical_solution" in path:
+        data["completion"] = data["canonical_solution"]
+    if evaluation_code==False:
+        test_case = data["test_case"]
+    else: 
+        test_case = data["small_test_cases"]
+    # test_case = data["small_test_cases"]
+    problem_idx = data["problem_idx"]
+    return_path,full_code = None,""
+    tmp_code = data["completion"].split("\n")
+    code = []
+    for string in tmp_code:
+        if "print(" in string:
+            continue
+        else:
+            code.append(string)
+    data["completion"] = "\n".join(code)
+    try:
+        if "class Solution" in data["completion"]:
+            if "```python" in data["completion"]:
+                start_idx = data["completion"].find("```python")
+                data["completion"] = data["completion"][start_idx+9:]
+                if "```" in data["completion"]:
+                    end_idx = data["completion"].find("```")
+                    data["completion"] = data["completion"][:end_idx]
+            test_case = test_case.split("\n")[:100]
+            test_case = "\n".join(test_case)
+            # import_pkg
+            full_code = import_pkg + "\n"+TreeNode_text + "\n"+ListNode_text + "\n" + data["completion"] + "\nsolution=Solution()\n" + test_case
+            # with open(f"./{path}/{problem_idx}.py", "w") as f:
+            #     f.write(full_code)
+            # return_path = f"./{path}/{problem_idx}.py"
+            result = check_correctness(full_code,timeout=10.0)
+            if result["passed"]:
+                with open(f"./{path}/{problem_idx}.py", "w") as f:
+                    f.write(full_code)
+                return_path = f"./{path}/{problem_idx}.py"
+                # print(return_path)
+            else:
+                return_path = None
+    except Exception as e:
+        # print(e)
+        pass
+    return return_path,full_code
+
+def calculate_code_execution_efficiency(data,evaluation_code=False,path="./tmp/",max_execution_time=10):
+    entry_point = ""
+    try:
+        if "task_id" in data.keys() and "HumanEval" in str(data["task_id"]):
+            problem_idx = data["task_id"].split("/")[1]
+            completion_file,full_code = humaneval_add_string_to_py_file(data,evaluation_code=evaluation_code, path=path)
+            entry_point = data["entry_point"]
+        # print(data.keys())
+        # print(data["dataset"])
+        elif "dataset" in data.keys() and data["dataset"]=="mbpp":
+            problem_idx = data["task_id"]
+            completion_file,full_code = mbpp_add_string_to_py_file(data,evaluation_code=evaluation_code, path=path)
+            code_example = data["code"]
+            match = re.search(r"def\s+(\w+)\s*\(", code_example)
+            if match:
+                entry_point = match.group(1)
+            else:
+                test_example = data["test_list"][0]
+                match = re.search(r"assert\s+(\w+)\s*\(", test_example)
+                if match:
+                    entry_point = match.group(1)
+                else: completion_file== None
+        else:
+            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)
+        completion_file = None
+    if completion_file == None:
+        # print("test")
+        overhead = f"""
+The code execution failed.
+"""
+        canonical_solution_memory_usage = 0
+        canonical_solution_execution_time = 0
+        canonical_solution_max_memory_usage = 0
+        executable = False
+        return overhead, canonical_solution_memory_usage, canonical_solution_execution_time, canonical_solution_max_memory_usage, executable
+
+    script_path = './run_code.sh'
+    completion_dat_file = f'./{path}/{problem_idx}.dat'
+    try:
+        subprocess.run([script_path, completion_file, completion_dat_file,str(max_execution_time)], 
+                            check=True, capture_output=True, text=True)
+        canonical_solution_memory_usage = calculate_memory_usage(completion_dat_file)
+        canonical_solution_execution_time = calculate_runtime(completion_dat_file)
+        canonical_solution_max_memory_usage = report_max_memory_usage(completion_dat_file)
+
+        executable = True
+        overhead = f"""
+The total memory usage during the code execution is: {canonical_solution_memory_usage} MB*s.
+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)
+        overhead = f"""
+The code execution failed.
+"""
+        canonical_solution_memory_usage = 0
+        canonical_solution_execution_time = 0
+        canonical_solution_max_memory_usage = 0
+        executable = False
+    return overhead, canonical_solution_memory_usage, canonical_solution_execution_time, canonical_solution_max_memory_usage, executable
+    
+    
+def fetch_completion(dataset,model):
+    with ThreadPoolExecutor() as executor:
+            future_to_entry = {executor.submit(calculate_code_execution_efficiency, copy.deepcopy(entry),False, path=model,max_execution_time=10): entry for entry in tqdm(dataset)}
+            for future in tqdm(concurrent.futures.as_completed(future_to_entry)):
+                entry = future_to_entry[future]
+                try:
+                    updated_entry = future.result()
+                    idx = dataset.index(entry)
+                    dataset[idx] = updated_entry
+                except Exception as e:
+                    print(e)
+    return dataset
+
+
+def run_model_task(model,task):
+
+    if "/" in model:
+        model = model.split("/")[1]
+    dat_path = f"./results/{task}_{model}"
+    canonical_solution_path = f"./results/{task}_canonical_solution"
+    with open(f"./{task}_{model}.json", "r") as f:
+        dataset = json.load(f)
+
+    if os.path.exists(dat_path):
+        shutil.rmtree(dat_path)
+    if os.path.exists(canonical_solution_path):
+        shutil.rmtree(canonical_solution_path)
+
+    if os.path.exists(dat_path) == False:
+        os.makedirs(dat_path)
+    if os.path.exists(canonical_solution_path) == False:
+        os.makedirs(canonical_solution_path)
+
+    fetch_completion(dataset,dat_path)
+
+    with open(f"./{task}_{model}.json", "r") as f:
+        dataset = json.load(f)
+    for i in range(len(dataset)):
+        dataset[i]["dataset"] = f"{task}"
+    fetch_completion(dataset,canonical_solution_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="gpt-4")
+    parser.add_argument("--task", type=str, default="EffiBench")
+    args = parser.parse_args()
+
+    run_model_task(args.model,args.task)

execution.py

@@ -0,0 +1,266 @@
+import contextlib
+import faulthandler
+import io
+import multiprocessing
+import os
+import platform
+import signal
+import random
+import subprocess
+import tempfile
+import gzip
+import json
+from typing import *
+
+def dicts_to_jsonl(data_list: list, filename: str, compress: bool = True) -> None:
+    """
+    Method saves list of dicts into jsonl file.
+    :param data: (list) list of dicts to be stored,
+    :param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends
+        .jsonl suffix into the file.
+    :param compress: (bool) should file be compressed into a gzip archive?
+    """
+    sjsonl = '.jsonl'
+    sgz = '.gz'
+    # Check filename
+    if not filename.endswith(sjsonl):
+        filename = filename + sjsonl
+    # Save data
+    
+    if compress:
+        filename = filename + sgz
+        with gzip.open(filename, 'w') as compressed:
+            for ddict in data_list:
+                jout = json.dumps(ddict) + '\n'
+                jout = jout.encode('utf-8')
+                compressed.write(jout)
+    else:
+        with open(filename, 'w') as out:
+            for ddict in data_list:
+                jout = json.dumps(ddict) + '\n'
+                out.write(jout)
+
+
+def check_correctness(
+        full_code: str,
+        timeout: float = 3.0,
+        tmp_dir: str = "tmp",
+) -> Dict:
+    """
+    Evaluates the functional correctness of a completion by running the test
+    suite provided in the problem.
+    """
+
+    def unsafe_execute(tmp_dir):
+        with create_tempdir():
+
+            # These system calls are needed when cleaning up tempdir.
+            import os
+            import shutil
+            rmtree = shutil.rmtree
+            rmdir = os.rmdir
+            chdir = os.chdir
+
+            # Disable functionalities that can make destructive changes to the test.
+            reliability_guard()
+
+            try:
+                exec_globals = {}
+                with swallow_io():
+                    with time_limit(timeout):
+                        exec(full_code, exec_globals)
+                    result.append("passed")
+            except TimeoutException:
+                result.append("timed out")
+            except AssertionError as e:
+                result.append(f"failed: AssertionError")
+            except BaseException as e:
+                result.append(f"failed: {e}")
+
+            # Needed for cleaning up.
+            shutil.rmtree = rmtree
+            os.rmdir = rmdir
+            os.chdir = chdir
+
+
+
+
+    manager = multiprocessing.Manager()
+    result = manager.list()
+
+    p = multiprocessing.Process(target=unsafe_execute, args=(tmp_dir,))
+    p.start()
+    p.join(timeout=timeout + 1)
+    if p.is_alive():
+        p.kill()
+
+    if not result:
+        result.append("timed out")
+
+    return {
+        "result"       : result[0],
+        "passed"       : result[0] == "passed",
+    }
+
+# Copyright (c) OpenAI (https://openai.com)
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+# ============================================================================
+@contextlib.contextmanager
+def time_limit(seconds: float):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+
+
+@contextlib.contextmanager
+def swallow_io():
+    stream = WriteOnlyStringIO()
+    with contextlib.redirect_stdout(stream):
+        with contextlib.redirect_stderr(stream):
+            with redirect_stdin(stream):
+                yield
+
+
+@contextlib.contextmanager
+def create_tempdir():
+    with tempfile.TemporaryDirectory() as dirname:
+        with chdir(dirname):
+            yield dirname
+
+
+class TimeoutException(Exception):
+    pass
+
+
+class WriteOnlyStringIO(io.StringIO):
+    """ StringIO that throws an exception when it's read from """
+
+    def read(self, *args, **kwargs):
+        raise IOError
+
+    def readline(self, *args, **kwargs):
+        raise IOError
+
+    def readlines(self, *args, **kwargs):
+        raise IOError
+
+    def readable(self, *args, **kwargs):
+        """ Returns True if the IO object can be read. """
+        return False
+
+
+class redirect_stdin(contextlib._RedirectStream):  # type: ignore
+    _stream = 'stdin'
+
+
+@contextlib.contextmanager
+def chdir(root):
+    if root == ".":
+        yield
+        return
+    cwd = os.getcwd()
+    os.chdir(root)
+    try:
+        yield
+    except BaseException as exc:
+        raise exc
+    finally:
+        os.chdir(cwd)
+
+
+def reliability_guard(maximum_memory_bytes: Optional[int] = None):
+    """
+    This disables various destructive functions and prevents the generated code
+    from interfering with the test (e.g. fork bomb, killing other processes,
+    removing filesystem files, etc.)
+
+    WARNING
+    This function is NOT a security sandbox. Untrusted code, including, model-
+    generated code, should not be blindly executed outside of one. See the 
+    Codex paper for more information about OpenAI's code sandbox, and proceed
+    with caution.
+    """
+
+    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))
+        if not platform.uname().system == 'Darwin':
+            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+
+    faulthandler.disable()
+
+    import builtins
+    builtins.exit = None
+    builtins.quit = None
+
+    import os
+    os.environ['OMP_NUM_THREADS'] = '1'
+
+    os.kill = None
+    os.system = None
+    os.putenv = None
+    os.remove = None
+    os.removedirs = None
+    os.rmdir = None
+    os.fchdir = None
+    os.setuid = None
+    os.fork = None
+    os.forkpty = None
+    os.killpg = None
+    os.rename = None
+    os.renames = None
+    os.truncate = None
+    os.replace = None
+    os.unlink = None
+    os.fchmod = None
+    os.fchown = None
+    os.chmod = None
+    os.chown = None
+    os.chroot = None
+    os.fchdir = None
+    os.lchflags = None
+    os.lchmod = None
+    os.lchown = None
+    os.getcwd = None
+    os.chdir = None
+
+    import shutil
+    shutil.rmtree = None
+    shutil.move = None
+    shutil.chown = None
+
+    import subprocess
+    subprocess.Popen = None  # type: ignore
+
+    __builtins__['help'] = None
+
+    import sys
+    sys.modules['ipdb'] = None
+    sys.modules['joblib'] = None
+    sys.modules['resource'] = None
+    sys.modules['psutil'] = None
+    sys.modules['tkinter'] = None

run_code.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+
+# Function to calculate memory usage in MB*seconds
+calculate_memory_usage() {
+    local dat_file=$1
+    awk 'BEGIN { 
+        prev_time = 0; 
+        prev_mem_mb = 0; 
+        mem_time_mb_s = 0;
+    }
+    NR>1 {
+        mem_in_mb = $2;
+        timestamp = $3;
+        
+        if (prev_time > 0) {
+            time_interval_s = timestamp - prev_time;
+            mem_time_mb_s += (prev_mem_mb + mem_in_mb) / 2 * time_interval_s;
+        }
+        
+        prev_time = timestamp;
+        prev_mem_mb = mem_in_mb;
+    }
+    END {
+        printf "%.2f\n", mem_time_mb_s;
+    }' "$dat_file"
+}
+
+
+completion_file="$1"
+completion_dat_file="$2"
+# file_type="$3"
+max_execution_time="$3"
+# Execute the specified file
+# max_execution_time=10
+echo "Executing $completion_file"
+error_output=$(mktemp)
+start_time=$(date +%s%N)
+rm -f "$completion_dat_file"
+timeout "$max_execution_time" mprof run --interval 0.001 --output "$completion_dat_file" "$completion_file" 2> "$error_output"
+end_time=$(date +%s%N)
+execution_time=$(( (end_time - start_time) / 1000000 ))
+exit_status=$?
+
+# Check execution status
+echo "Execution status: $exit_status"
+if [ $exit_status -ne 0 ] || [ -s "$error_output" ]; then
+    echo "Execution failed or errors were reported for $completion_file. Removing .dat file."
+    rm -f "$completion_dat_file"
+elif [ -f "$completion_dat_file" ]; then
+    mem_usage_mb_s=$(calculate_memory_usage "$completion_dat_file")
+    echo "Memory usage (MB*seconds): $mem_usage_mb_s"
+else
+    echo "Execution completed but no .dat file found for $completion_file."
+fi
+
+# Clean up
+rm -f "$error_output"