[tab] evaluation/evaluate_rewrites.py (python)
"""
Script to evaluate the quality of the rewrites of the thought traces.

For now we have three different evaluation methods:
1. Rewriting correctness: Does rewriting make sense → prompt powerful model to check rewritten traces
2. Rewriting length/difficulty correlation → are easy problems being rewritten shorter etc → Check token length
3. Check answer correctness → Regex
"""

import re
from openai import OpenAI
import os


openai = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
)

# ========= Check rewriting correctness =========
def check_rewriting_correctness(original, rewritten):
    """
    Checks if the rewriting makes sense.
    """
    prompt = f"""
    You are a helpful assistant that checks if the rewriting makes sense.
    The original thought trace is: {original}
    The rewritten thought trace is: {rewritten}
    Return True if the rewriting makes sense, False otherwise. Go through the rewritten trace step by step and think about whether the steps still make sense or not.
    For your final verdict, follow the following format:
    [Final Verdict: True/False]
    [Reasoning: 2-3 sentences explaining your verdict]
    """
    response = openai.chat.completions.create(
        model="gpt-5",
        messages=[{"role": "system", "content": "You are a helpful assistant that checks if the rewriting makes sense."}, {"role": "user", "content": prompt}]
    )
    if "Final Verdict: True" in response.choices[0].message.content:
        return True
    elif "Final Verdict: False" in response.choices[0].message.content:
        return False
    else:
        raise ValueError(f"Invalid response: {response.choices[0].message.content}")

# ========= Check rewriting length/difficulty correlation =========
def check_rewriting_length_difficulty_correlation(original, rewritten, difficulty):
    """
    Checks if the rewriting length is correlated with the difficulty of the problem.
    """
    if difficulty == "Easy":
        return len(original) < len(rewritten)
    elif difficulty == "Hard":
        return True
    else:
        raise ValueError(f"Invalid difficulty: {difficulty}")


# ========= Check answer correctness =========
def extract_boxed_answer(s):
    """
    Extracts the content within \boxed{} from a string.
    Returns None if no boxed answer is found.
    """
    match = re.search(r'\\boxed\s*{([^}]*)}', s)
    if match:
        return match.group(1).strip()
    return None

def check_boxed_answer_correctness(original, rewritten):
    """
    Checks if both original and rewritten answers contain a boxed answer,
    and returns True if they are the same, False if they differ,
    and None if a boxed answer is not found in either.
    """
    orig_boxed = extract_boxed_answer(original)
    rew_boxed = extract_boxed_answer(rewritten)
    if orig_boxed is not None and rew_boxed is not None:
        return orig_boxed == rew_boxed
    elif orig_boxed is None and rew_boxed is None:
        return False
    else:
        return None


def evaluate_rewrites(original_list, rewritten_list, difficulty_list):

    rewrite_correctness_list = []
    rewrite_length_difficulty_correlation_list = []
    boxed_answer_correctness_list = []
    for original, rewritten, difficulty in zip(original_list, rewritten_list, difficulty_list):
        rewrite_correctness_list.append(check_rewriting_correctness(original, rewritten))
        rewrite_length_difficulty_correlation_list.append(check_rewriting_length_difficulty_correlation(original, rewritten, difficulty))
        boxed_answer_correctness_list.append(check_boxed_answer_correctness(original, rewritten))

    rewrite_correctness_score = sum(rewrite_correctness_list) / len(rewrite_correctness_list)
    rewrite_length_difficulty_correlation_score = sum(rewrite_length_difficulty_correlation_list) / len(rewrite_length_difficulty_correlation_list)
    boxed_answer_correctness_score = sum(boxed_answer_correctness_list) / len(boxed_answer_correctness_list)

    return rewrite_correctness_score, rewrite_length_difficulty_correlation_score, boxed_answer_correctness_score

if __name__ == "__main__":
    original = "Let's start by defining the variables. Let A = 5 and B = 10. Then, we can calculate the sum of A and B: A + B = 5 + 10 = 15. Therefore, the answer is \\boxed{15}."
    rewritten = "If A = 5 and B = 10, then A + B = 5 + 10 = 15. so, answer is \\boxed{15}."
    difficulty = "Hard"
    rewrite_correctness_score, rewrite_length_difficulty_correlation_score, boxed_answer_correctness_score = evaluate_rewrites([original], [rewritten], [difficulty])
    print(f"Rewrite correctness score: {rewrite_correctness_score * 100}%")
    print(f"Rewrite length/difficulty correlation score: {rewrite_length_difficulty_correlation_score * 100}%")
    print(f"Boxed answer correctness score: {boxed_answer_correctness_score * 100}%")