PyPy3 execution support for LiveCodeBench evaluation

nemo_skills/evaluation/evaluator/code.py

@@ -15,14 +15,17 @@
 
 import json
 import logging
+import os
 import re
 import shutil
 import subprocess
 import sys
 from argparse import Namespace
+from dataclasses import field
 
 from omegaconf import OmegaConf
 
+from nemo_skills.code_execution.sandbox import get_sandbox
 from nemo_skills.utils import get_logger_name, nested_dataclass, unroll_files
 
 LOG = logging.getLogger(get_logger_name(__file__))
@@ -86,29 +89,29 @@ def install_from_git(git_url):
         print(f"Error during installation: {e}")
 
 
-# TODO: use sandbox
 @nested_dataclass(kw_only=True)
 class LiveCodeBenchEvaluatorConfig:
+    sandbox: dict = field(default_factory=lambda: {'sandbox_type': 'local'})
     language: str = "python"  # "cpp" is another option now
     test_file: str = None
+    timeout: float = 30.0
 
 
 def eval_livecodebench(cfg):
-    try:
-        from livecodebench.evaluate import evaluate
-    except ImportError:
-        LOG.info("Package 'livecodebench' not found. Attempting to install...")
-        install_from_git("git+https://github.com/wasiahmad/livecodebench.git")
-        try:
-            from livecodebench.evaluate import evaluate
-        except ImportError:
-            LOG.info("Failed to install 'livecodebench'. Please install it manually.")
-            raise
+    # try:
+    #     from livecodebench.evaluate import evaluate
+    # except ImportError:
+    #     LOG.info("Package 'livecodebench' not found. Attempting to install...")
+    #     install_from_git("git+https://github.com/wasiahmad/livecodebench.git")
+    #     try:
+    #         from livecodebench.evaluate import evaluate
+    #     except ImportError:
+    #         LOG.info("Failed to install 'livecodebench'. Please install it manually.")
+    #         raise
 
     eval_config = LiveCodeBenchEvaluatorConfig(_init_nested=True, **cfg.eval_config)
     assert eval_config.language in ["python", "cpp"]
-    if eval_config.language == "cpp":
-        assert eval_config.test_file is not None
+    assert eval_config.test_file is not None
 
     release_version = None
     for jsonl_file in unroll_files(cfg.input_files):
@@ -129,16 +132,44 @@ def eval_livecodebench(cfg):
             for sample in samples:
                 f.write(json.dumps(sample) + "\n")
 
-        # https://github.com/wasiahmad/livecodebench/blob/main/livecodebench/evaluate.py#L10
-        evaluate(
-            custom_output_file=jsonl_file,
-            release_version=f"release_{release_version}",
-            k_list=[1],
-            language=eval_config.language,
-            test_file=None if eval_config.language == "python" else eval_config.test_file,
-            num_process_evaluate=12,
-            timeout=6 if eval_config.language == "python" else 30,
-        )
+        base_dir = os.path.dirname(os.path.abspath(__file__))
+        script_to_run = os.path.join(base_dir, 'livecodebench', 'evaluate.py')
+
+        code = f"""
+import subprocess
+
+command = [
+    'pypy3',
+    '{script_to_run}',
+    '--custom_output_file',
+    '{jsonl_file}',
+    '--release_version',
+    'release_{release_version}',
+    '--k_list',
+    '1',
+    '--language',
+    '{eval_config.language}',
+    '--test_file',
+    '{eval_config.test_file}',
+    '--num_process_evaluate',
+    '12',
+     '--timeout',
+    '{eval_config.timeout}'
+]
+try:
+    result = subprocess.run(
+        command,
+        capture_output=True,
+        text=True,
+        check=True
+    )
+    print(result.stdout)
+except subprocess.CalledProcessError as e:
+    print("--- Stderr ---")
+    print(e.stderr)
+"""
+        sandbox = get_sandbox(**eval_config.sandbox)
+        output_dict, _ = sandbox.execute_code(code, timeout=eval_config.timeout, max_output_characters=100000)
 
         with open(jsonl_file[:-6] + '_eval_results.json', 'rt', encoding="utf-8") as fin:
             eval_grades = json.load(fin)

nemo_skills/evaluation/evaluator/livecodebench/code_generation.py

@@ -0,0 +1,129 @@
+import base64
+import json
+import pickle
+import zlib
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+
+
+class Platform(Enum):
+    LEETCODE = "leetcode"
+    CODEFORCES = "codeforces"
+    ATCODER = "atcoder"
+
+
+class Difficulty(Enum):
+    EASY = "easy"
+    MEDIUM = "medium"
+    HARD = "hard"
+
+
+class TestType(Enum):
+    STDIN = "stdin"
+    FUNCTIONAL = "functional"
+
+
+@dataclass
+class Test:
+    input: str
+    output: str
+    testtype: TestType
+
+    def __post_init__(self):
+        self.testtype = TestType(self.testtype)
+        # if self.testtype == TestType.FUNCTIONAL:
+        #     self.input = json.loads(self.input)
+        #     self.output = json.loads(self.output)
+
+
+@dataclass
+class CodeGenerationProblem:
+    question_title: str
+    question_content: str
+    platform: Platform
+    question_id: str
+    contest_id: str
+    contest_date: datetime
+    starter_code: str
+    difficulty: Difficulty
+    public_test_cases: list[Test]
+    private_test_cases: list[Test]
+    metadata: dict
+    language: str
+
+    def __post_init__(self):
+        self.platform = Platform(self.platform)
+        self.difficulty = Difficulty(self.difficulty)
+        self.contest_date = datetime.fromisoformat(self.contest_date)
+
+        if 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]
+        else:
+            self.public_test_cases = []
+
+        try:
+            self.private_test_cases = json.loads(self.private_test_cases)  # type: ignore
+        except:
+            self.private_test_cases = json.loads(
+                pickle.loads(
+                    zlib.decompress(base64.b64decode(self.private_test_cases.encode("utf-8")))  # type: ignore
+                )
+            )  # type: ignore
+        self.private_test_cases = [Test(**t) for t in self.private_test_cases]
+
+        self.metadata = json.loads(self.metadata)  # type: ignore
+
+    def insert_output(self, output_list: list[str], code_list: list[str]) -> dict:
+        return {
+            "question_title": self.question_title,
+            "question_content": self.question_content,
+            "platform": self.platform.value,
+            "question_id": self.question_id,
+            "contest_id": self.contest_id,
+            "contest_date": self.contest_date.isoformat(),
+            "starter_code": self.starter_code,
+            "difficulty": self.difficulty.value,
+            "output_list": output_list,
+            "code_list": code_list,
+            "language": self.language,
+        }
+
+    def insert_output_evaluation(
+        self,
+        output_list: list[str],
+        code_list: list[str],
+        graded_list: list[bool],
+        **kwargs,
+    ) -> dict:
+        output = self.insert_output(output_list, code_list)
+        output["graded_list"] = graded_list
+        output["pass@1"] = graded_list.count(True) / len(graded_list)
+        for k, v in kwargs.items():
+            output[k] = v
+        return output
+
+    def get_evaluation_sample(self):
+        return {
+            "input_output": json.dumps(
+                {
+                    "inputs": [t.input for t in self.public_test_cases + self.private_test_cases],
+                    "outputs": [t.output for t in self.public_test_cases + self.private_test_cases],
+                    "fn_name": self.metadata.get("func_name", None),
+                }
+            ),
+        }
+
+
+def load_code_generation_dataset_from_file(filepath, language) -> list[CodeGenerationProblem]:
+    dataset = []
+    with open(filepath, "r") as f:
+        for line in f:
+            p = json.loads(line)
+            if "task_id" in p:
+                assert "question_id" not in p
+                p["question_id"] = p.pop("task_id")
+            dataset.append(CodeGenerationProblem(**p, language=language))
+    print(f"Loaded {len(dataset)} problems")
+    return dataset

nemo_skills/evaluation/evaluator/livecodebench/evaluate.py

@@ -0,0 +1,97 @@
+import json
+from datetime import datetime
+
+from nemo_skills.evaluation.evaluator.livecodebench.code_generation import load_code_generation_dataset_from_file
+from nemo_skills.evaluation.evaluator.livecodebench.metrics import codegen_metrics
+from nemo_skills.evaluation.evaluator.livecodebench.pass_k_utils import extract_instance_results
+
+
+def evaluate(
+    custom_output_file: str,
+    release_version: str = "release_latest",
+    k_list=[1],
+    language: str = "python",
+    test_file: str = None,
+    num_process_evaluate: int = 12,
+    debug: bool = False,
+    timeout: int = 6,
+):
+    assert test_file is not None
+    custom_outputs = dict()
+    with open(custom_output_file, "r") as f:
+        for line in f:
+            output = json.loads(line)
+            if "question_id" not in output:
+                assert "task_id" in output
+                output["question_id"] = output.pop("task_id")
+            custom_outputs[output["question_id"]] = output
+
+    benchmark = load_code_generation_dataset_from_file(test_file, language)
+    benchmark = [problem for problem in benchmark if problem.question_id in custom_outputs]
+    assert len(custom_outputs) == len(benchmark), f"{len(custom_outputs)} != {len(benchmark)}"
+    assert all(isinstance(custom_output, dict) for custom_output in custom_outputs.values())
+
+    save_results, combined_results = [], []
+    for instance in benchmark:
+        code_list = custom_outputs[instance.question_id]["code_list"]
+        output = instance.insert_output(code_list, code_list)
+        save_results.append(output)
+        combined_results.append((code_list, code_list))
+
+    eval_samples = [instance.get_evaluation_sample() for instance in benchmark]
+    generations = [extracted for _, extracted in combined_results]
+
+    metrics = codegen_metrics(
+        eval_samples,
+        generations,
+        k_list=k_list,
+        num_process_evaluate=num_process_evaluate,
+        timeout=timeout,
+        debug=debug,
+        language=language,
+    )
+
+    graded = extract_instance_results(metrics[1])
+
+    metadatas = metrics[2]
+    save_eval_results = [
+        instance.insert_output_evaluation(outputs_list, extracted_list, graded_list, metadata=meta)
+        for instance, (outputs_list, extracted_list), graded_list, meta in zip(
+            benchmark, combined_results, graded, metadatas
+        )
+    ]
+
+    # save_eval_results
+    output_results = dict()
+    output_results["date"] = datetime.now().strftime("%Y-%m-%d %H:%M")
+    for k in metrics[0]:
+        if k.startswith("pass@"):
+            print(f"{k}: {metrics[0][k]}")
+            output_results[k] = metrics[0][k]
+
+    output_results["detail_pass@1"] = dict()
+    output_results["eval"] = dict()
+    difficulty_wise_pass_at_1 = dict()
+    for r in save_eval_results:
+        output_results["eval"][r["question_id"]] = r
+        if r["difficulty"] not in difficulty_wise_pass_at_1:
+            difficulty_wise_pass_at_1[r["difficulty"]] = []
+        difficulty_wise_pass_at_1[r["difficulty"]].append(r["pass@1"])
+
+    for tag, v in difficulty_wise_pass_at_1.items():
+        pass_at_1 = sum(v) / len(v)
+        print(f"{tag} pass@1: {pass_at_1}")
+        output_results["detail_pass@1"][tag] = pass_at_1
+
+    with open(custom_output_file[:-6] + "_eval_results.json", "w") as f:
+        json.dump(output_results, f, indent=4)
+
+
+def main():
+    from fire import Fire
+
+    Fire(evaluate)
+
+
+if __name__ == "__main__":
+    main()