Naman-ntc · Apr 14, 2024 · Apr 11, 2024 · Apr 14, 2024
diff --git a/lcb_runner/benchmarks/__init__.py b/lcb_runner/benchmarks/__init__.py
@@ -6,3 +6,7 @@
     TestOutputPredictionProblem,
     load_test_prediction_dataset,
 )
+from lcb_runner.benchmarks.code_execution import (
+    CodeExecutionProblem,
+    load_code_execution_dataset,
+)
diff --git a/lcb_runner/benchmarks/code_execution.py b/lcb_runner/benchmarks/code_execution.py
@@ -0,0 +1,68 @@
+import json
+from enum import Enum
+from datetime import datetime
+from dataclasses import dataclass
+
+from datasets import load_dataset
+
+@dataclass
+class CodeExecutionProblem:
+    question_id: str
+    contest_id: str
+    contest_date: datetime
+    difficulty: str
+    function_name: str
+    code: str
+    input: str
+    output: str
+    id: str
+    problem_id: str
+    numsteps: int
+
+
+    def __post_init__(self):
+        pass
+
+    def insert_output(self, output_list: list[str], pred_list: list[str]) -> dict:
+        return {
+            "question_id": self.question_id,
+            "contest_id": self.contest_id,
+            "contest_date": self.contest_date.isoformat(),
+            "difficulty": self.difficulty,
+            "function_name": self.function_name,
+            "code": self.code,
+            "input": self.input,
+            "output": self.output,
+            "id": self.id,
+            "problem_id": self.problem_id,
+            "numsteps": self.numsteps,
+            "output_list": output_list,
+            "pred_list": pred_list,
+        }
+
+    def insert_output_evaluation(
+        self, output_list: list[str], code_list: list[str], graded_list: list[bool]
+    ) -> dict:
+        output = self.insert_output(output_list, code_list)
+        output["graded_list"] = graded_list
+        output["pass@1"] = graded_list.count(True) / len(graded_list)
+        return output
+
+    def get_evaluation_sample(self) -> dict:
+        return {
+            "code": self.code,
+            "input": self.input,
+            "output": self.output,
+        }
+
+
+
+def load_code_execution_dataset() -> list[CodeExecutionProblem]:
+    dataset = load_dataset("livecodebench/execution-v2", split="test")
+    dataset = [CodeExecutionProblem(**p) for p in dataset]
+    print(f"Loaded {len(dataset)} problems")
+    return dataset
+
+
+if __name__ == "__main__":
+    dataset = load_code_execution_dataset()
diff --git a/lcb_runner/benchmarks/code_generation.py b/lcb_runner/benchmarks/code_generation.py
@@ -54,13 +54,13 @@ def __post_init__(self):
         self.platform = Platform(self.platform)
         self.difficulty = Difficulty(self.difficulty)
 
-        self.public_test_cases = json.loads(self.public_test_cases)
+        self.public_test_cases = json.loads(self.public_test_cases)  # type: ignore
         self.public_test_cases = [Test(**t) for t in self.public_test_cases]
 
-        self.private_test_cases = json.loads(self.private_test_cases)
+        self.private_test_cases = json.loads(self.private_test_cases)  # type: ignore
         self.private_test_cases = [Test(**t) for t in self.private_test_cases]
 
-        self.metadata = json.loads(self.metadata)
+        self.metadata = json.loads(self.metadata)  # type: ignore
 
     def insert_output(self, output_list: list[str], code_list: list[str]) -> dict:
         return {
@@ -104,7 +104,7 @@ def get_evaluation_sample(self):
 
 def load_code_generation_dataset() -> list[CodeGenerationProblem]:
     dataset = load_dataset("livecodebench/code_generation", split="test")
-    dataset = [CodeGenerationProblem(**p) for p in dataset]
+    dataset = [CodeGenerationProblem(**p) for p in dataset]  # type: ignore
     print(f"Loaded {len(dataset)} problems")
     return dataset
 

diff --git a/lcb_runner/benchmarks/test_output_prediction.py b/lcb_runner/benchmarks/test_output_prediction.py
@@ -27,7 +27,7 @@ class TestOutputPredictionProblem:
     test_id: int
 
     def __post_init__(self):
-        self.test = [Test(**t) for t in json.loads(self.test)]
+        self.test = [Test(**t) for t in json.loads(self.test)]  # type: ignore
 
     def insert_output(self, output_list: list[str], pred_list: list[str]) -> dict:
         return {
@@ -60,8 +60,7 @@ def get_evaluation_sample(self) -> dict:
 
 
 def load_test_prediction_dataset() -> list[TestOutputPredictionProblem]:
-    dataset = load_dataset("livecodebench/test_generation", split="test")
-    print(dataset[0].keys())
+    dataset = load_dataset("livecodebench/test_generation", split="test")  # type: ignore
     dataset = [TestOutputPredictionProblem(**d) for d in dataset]
     print(f"Loaded {len(dataset)} prediction problems")
     return dataset

diff --git a/lcb_runner/evaluation/__init__.py b/lcb_runner/evaluation/__init__.py
@@ -1,4 +1,5 @@
 from lcb_runner.evaluation.compute_code_generation_metrics import codegen_metrics
+from lcb_runner.evaluation.compute_code_execution_metrics import code_execution_metrics
 from lcb_runner.evaluation.compute_test_output_prediction_metrics import (
     test_output_metrics,
 )

diff --git a/lcb_runner/evaluation/compute_code_execution_metrics.py b/lcb_runner/evaluation/compute_code_execution_metrics.py
@@ -0,0 +1,56 @@
+import numpy as np
+from concurrent.futures import ProcessPoolExecutor
+import tqdm
+
+from lcb_runner.evaluation.utils_execute import BASE_IMPORTS, check_correctness
+
+def evaluate_score(args) -> list[bool]:
+    gs, (c, i, o) = args
+
+    execution_results = []
+    for g in gs:
+        if i in g:
+            pass
+        else:
+            code_to_execute = f"{BASE_IMPORTS}\n{c}\nassert {o} == {g}"
+            execution_results.append(check_correctness(code_to_execute, 3))
+    if len(execution_results) == 0:
+        execution_results = [False] * len(gs)
+    return execution_results
+
+def pass_at_k(n, c, k):
+    if n - c < k: return 1.0
+    return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
+
+def code_execution_metrics(
+    samples,
+    generations,
+):
+    # execute the code
+    references = [(doc["code"], doc["input"], doc["output"]) for doc in samples]
+    with ProcessPoolExecutor() as executor:
+        args_list = zip(generations, references)
+        results = executor.map(evaluate_score, args_list)
+    all_results = list(results)
+
+    # serial version
+    # all_results = []
+    # for i in range(len(generations)):
+    #     generation = generations[i]
+    #     result = evaluate_score([generation, references[i]])
+    #     all_results.append(result)
+
+    # compute pass@1
+    pass_at_1s = []
+    for execution_result in all_results:
+        c, n = execution_result.count(True), len(execution_result)
+        pass_at_1s.append(pass_at_k(n, c, 1))
+    metrics = {"pass@1": sum(pass_at_1s) / len(pass_at_1s) * 100}
+
+    results = {}
+    for i, r in enumerate(all_results):
+        r_new = []
+        for _r in r:
+            r_new.append([_r])
+        results[i] = r_new
+    return metrics, results