Function compile_inline_reward

backend/python/trl/reward_functions.py:133–175

Compile user-provided code into a reward function. The code should be the body of a function that receives `completions` (list[str]) and `**kwargs`, and returns list[float]. Available modules: re, math, json, string.

(name, code)

def compile_inline_reward(name, code):
    """Compile user-provided code into a reward function.

    The code should be the body of a function that receives
    `completions` (list[str]) and `**kwargs`, and returns list[float].

    Available modules: re, math, json, string.
    """
    func_source = (
        f"def _user_reward_{name}(completions, **kwargs):\n"
        + "\n".join(f"    {line}" for line in code.splitlines())
    )

    restricted_globals = {
        "__builtins__": _SAFE_BUILTINS,
        "re": re,
        "math": math,
        "json": json,
        "string": string,
    }

    try:
        compiled = compile(func_source, f"<inline-reward-{name}>", "exec")
    except SyntaxError as e:
        raise ValueError(f"Syntax error in inline reward function '{name}': {e}")

    exec(compiled, restricted_globals)
    func = restricted_globals[f"_user_reward_{name}"]

    # Validate with a quick smoke test
    try:
        result = func(["test"], answer=["test"])
        if not isinstance(result, list):
            raise ValueError(
                f"Inline reward function '{name}' must return a list, got {type(result).__name__}"
            )
    except Exception as e:
        if "must return a list" in str(e):
            raise
        # Other errors during smoke test are acceptable (e.g. missing kwargs)
        pass

    return func


# ---------------------------------------------------------------------------