src/rewards.py

import ast
import math
import operator
import re
from collections import Counter


ALLOWED_BINOPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
}


def safe_eval_expression(expression):
    tree = ast.parse(expression, mode="eval")
    used = []

    def visit(node):
        if isinstance(node, ast.Expression):
            return visit(node.body)
        if isinstance(node, ast.Constant) and isinstance(node.value, int) and node.value >= 0:
            used.append(node.value)
            return node.value
        if isinstance(node, ast.BinOp) and type(node.op) in ALLOWED_BINOPS:
            left, right = visit(node.left), visit(node.right)
            if isinstance(node.op, ast.Div) and right == 0:
                raise ValueError("division by zero")
            return ALLOWED_BINOPS[type(node.op)](left, right)
        raise ValueError(f"unsupported expression node: {type(node).__name__}")

    value = float(visit(tree))
    if not math.isfinite(value):
        raise ValueError("non-finite result")
    return value, used


def completion_to_text(completion):
    if isinstance(completion, str):
        return completion
    if isinstance(completion, list):
        return "\n".join(
            str(item.get("content", "")) if isinstance(item, dict) else str(item)
            for item in completion
        )
    if isinstance(completion, dict):
        return str(completion.get("content", completion))
    return str(completion)


def extract_answer_expression(text):
    match = re.search(r"<answer>\s*(.*?)\s*</answer>", text, flags=re.I | re.S)
    return match.group(1).strip() if match else None


def expression_candidates(text):
    candidates = []
    tagged = extract_answer_expression(text)
    if tagged:
        candidates.append(tagged)
    boxed = re.findall(r"\\boxed\{([^{}]+)\}", text)
    candidates.extend(reversed(boxed))
    for line in reversed(text.splitlines()):
        line = line.strip().strip("`$ ")
        line = re.sub(r"^(final answer|answer|expression)\s*[:=]\s*", "", line, flags=re.I)
        line = re.sub(r"\s*=\s*-?\d+(?:\.\d+)?\s*$", "", line)
        if re.fullmatch(r"[\d\s()+\-*/.]+", line) and re.search(r"[+\-*/]", line):
            candidates.append(line)
    seen = set()
    return [candidate for candidate in candidates if not (candidate in seen or seen.add(candidate))]


def split_answer_equation(candidate):
    if "=" not in candidate:
        return candidate.strip(), None
    if candidate.count("=") != 1:
        raise ValueError("only one equality is allowed")
    expression, claimed_result = (part.strip() for part in candidate.split("=", 1))
    if not expression or not re.fullmatch(r"[-+]?\d+(?:\.\d+)?", claimed_result):
        raise ValueError("equality must end in a numeric result")
    return expression, float(claimed_result)


def score_countdown(text, numbers, target):
    candidates = expression_candidates(text)
    if not candidates:
        return {
            "correct": 0.0,
            "format": 0.0,
            "valid_numbers": 0.0,
            "proximity": 0.0,
            "expression": None,
        }
    best = None
    for candidate in candidates:
        try:
            expression, claimed_result = split_answer_equation(candidate)
            value, used = safe_eval_expression(expression)
            if claimed_result is not None and abs(value - claimed_result) >= 1e-6:
                continue
        except Exception:
            continue
        valid_numbers = Counter(used) == Counter(int(number) for number in numbers)
        error = abs(value - float(target))
        result = {
            "correct": float(valid_numbers and error < 1e-6),
            "format": 0.1 if extract_answer_expression(text) else 0.0,
            "valid_numbers": 0.2 if valid_numbers else 0.0,
            "proximity": (0.2 / (1.0 + error)) if valid_numbers else 0.0,
            "expression": candidate,
        }
        if best is None or (result["correct"], result["valid_numbers"], result["proximity"]) > (
            best["correct"], best["valid_numbers"], best["proximity"]
        ):
            best = result
    return best or {
        "correct": 0.0,
        "format": 0.0,
        "valid_numbers": 0.0,
        "proximity": 0.0,
        "expression": candidates[0],
    }


def countdown_reward(completions, numbers, target, **kwargs):
    return [
        score_countdown(completion_to_text(completion), nums, tgt)["correct"]
        for completion, nums, tgt in zip(completions, numbers, target)
    ]


def format_reward(completions, numbers, target, **kwargs):
    return [
        score_countdown(completion_to_text(completion), nums, tgt)["format"]
        for completion, nums, tgt in zip(completions, numbers, target)
    ]


def valid_numbers_reward(completions, numbers, target, **kwargs):
    return [
        score_countdown(completion_to_text(completion), nums, tgt)["valid_numbers"]
        for completion, nums, tgt in zip(completions, numbers, target)
    ]


def proximity_reward(completions, numbers, target, **kwargs):
    return [
        score_countdown(completion_to_text(completion), nums, tgt)["proximity"]
        for completion, nums, tgt in zip(completions, numbers, target)
    ]