|
|
from typing import Dict, List |
|
|
import re |
|
|
|
|
|
from larm.data.envs.base_env import StaticEnv |
|
|
from larm.common.registry import registry |
|
|
from larm.data.utils.code_utils import PyExecutor, extract_python_code |
|
|
|
|
|
|
|
|
|
|
|
@registry.register_env("kodcode") |
|
|
class KodCodeEnv(StaticEnv): |
|
|
|
|
|
def __init__(self, config): |
|
|
super().__init__(config) |
|
|
|
|
|
@classmethod |
|
|
def _rename_func(cls, answer: str, function_name: str) -> str: |
|
|
""" |
|
|
Replace the name of the first function in `answer` with `function_name`. |
|
|
Only modifies the function name, keeps everything else intact. |
|
|
""" |
|
|
pattern = r"def\s+(\w+)\s*\(" |
|
|
|
|
|
new_answer = re.sub(pattern, f"def {function_name}(", answer, count=1) |
|
|
return new_answer |
|
|
|
|
|
@classmethod |
|
|
def _accuracy_reward(cls, completions: List[str], test: List[str], test_info, **kwargs) -> List[float]: |
|
|
|
|
|
py_executor = PyExecutor() |
|
|
scores = [] |
|
|
for completion, t, tf in zip(completions, test, test_info): |
|
|
func_blocks = extract_python_code(completion.strip()) |
|
|
collected_answer = '\n'.join(func_blocks) |
|
|
renamed_answer = cls._rename_func(collected_answer, tf[0]["function_name"]) |
|
|
_, _, results = py_executor.execute(renamed_answer, [t]) |
|
|
|
|
|
score = sum(results) / len(results) |
|
|
scores.append(score) |
|
|
|
|
|
return scores |
|
|
|
|
|
@classmethod |
|
|
def _format_reward(cls, completions: List[str], **kwargs): |
|
|
pass |
|
|
|
