GSM8K › Code Implementation: Evaluating Models on GSM8K

import subprocess

subprocess.run(["uv", "pip", "install", "datasets"], check=True)

import random

from datasets import load_dataset

# Set seed for reproducibility
random.seed(42)

# Load GSM8K dataset from HuggingFace
dataset = load_dataset("openai/gsm8k", "main")
train_data = dataset["train"]
test_data = dataset["test"]

# Examine a sample problem
sample = test_data[0]

import re


def extract_answer_from_text(text: str) -> str:
    """
    Extract the final numeric answer from GSM8K formatted text.
    Looks for #### delimiter first, then falls back to final number.
    """
    # First try to find #### delimiter
    if "####" in text:
        match = re.search(r"####\s*(-?\d+)", text)
        if match:
            return match.group(1)

    # Fallback: find last number in text
    numbers = re.findall(r"-?\d+", text.replace(",", ""))
    if numbers:
        return numbers[-1]

    return ""


def normalize_answer(answer: str) -> str:
    """Normalize answer by removing commas and whitespace."""
    return answer.replace(",", "").strip()


# Test extraction on our sample
ground_truth = sample["answer"]
extracted = extract_answer_from_text(ground_truth)

import random
from typing import Dict, List


def create_cot_prompt(
    examples: List[Dict], test_question: str, num_shots: int = 5
) -> str:
    """
    Create a chain-of-thought prompt with few-shot examples.
    Each example shows the question and full step-by-step answer.
    """
    selected_examples = random.sample(examples, num_shots)

    prompt_parts = []
    for ex in selected_examples:
        prompt_parts.append(f"Q: {ex['question']}")
        prompt_parts.append(f"A: {ex['answer']}\n")

    prompt_parts.append(f"Q: {test_question}")
    prompt_parts.append("A:")

    return "\n".join(prompt_parts)


# Create a CoT prompt for our sample question
cot_prompt = create_cot_prompt(
    [train_data[i] for i in range(10)],  # Use first 10 for the pool
    sample["question"],
    num_shots=3,
)

def simulate_model_response(question: str, ground_truth_answer: str) -> str:
    """
    Simulate a model response for demonstration purposes.
    In practice, this would call an LLM API.
    Produces correct answers 70% of the time with proper chain-of-thought,
    and off-by-small arithmetic errors the remaining 30%.
    """
    extracted_gt = extract_answer_from_text(ground_truth_answer)

    # 70% chance of correct answer with proper CoT
    if random.random() < 0.7:
        reasoning = ground_truth_answer.split("####")[0].strip()
        return f"{reasoning}\n#### {extracted_gt}"
    else:
        # Simulate arithmetic error: answer is off by a small amount
        wrong_answer = str(int(extracted_gt) + random.choice([-1, 1, -2, 2]))
        reasoning = ground_truth_answer.split("####")[0].strip()
        return f"{reasoning}\n#### {wrong_answer}"


def evaluate_gsm8k_sample(
    question: str, ground_truth: str, model_response: str
) -> tuple[bool, str, str]:
    """
    Evaluate a single GSM8K sample.
    Returns: (is_correct, predicted_answer, ground_truth_answer)
    """
    pred = normalize_answer(extract_answer_from_text(model_response))
    true = normalize_answer(extract_answer_from_text(ground_truth))

    return pred == true, pred, true


# Test evaluation on our sample
mock_response = simulate_model_response(sample["question"], sample["answer"])
is_correct, pred, true = evaluate_gsm8k_sample(
    sample["question"], sample["answer"], mock_response
)

def evaluate_gsm8k_batch(
    test_samples: List[Dict],
    train_pool: List[Dict],
    num_shots: int = 5,
    sample_size: int = 50,
) -> Dict:
    """
    Evaluate a batch of GSM8K problems with chain-of-thought prompting.
    Returns a dict with total, correct, accuracy, and per-sample predictions.
    """
    results = {"total": 0, "correct": 0, "predictions": []}

    test_subset = test_samples[:sample_size]

    for item in test_subset:
        # Create prompt (in practice, would call LLM here)
        prompt = create_cot_prompt(train_pool, item["question"], num_shots)

        # Simulate model response
        response = simulate_model_response(item["question"], item["answer"])

        # Evaluate
        is_correct, pred, true = evaluate_gsm8k_sample(
            item["question"], item["answer"], response
        )

        results["total"] += 1
        if is_correct:
            results["correct"] += 1

        results["predictions"].append(
            {
                "question": item["question"][:100] + "...",
                "predicted": pred,
                "ground_truth": true,
                "correct": is_correct,
            }
        )

    results["accuracy"] = (
        results["correct"] / results["total"] if results["total"] > 0 else 0
    )
    return results


# Run evaluation on a small subset
train_pool = [train_data[i] for i in range(100)]
test_list = list(test_data)

evaluation_results = evaluate_gsm8k_batch(
    test_list, train_pool, num_shots=3, sample_size=20
)

from collections import Counter


def self_consistency_evaluate(
    question: str,
    ground_truth: str,
    train_pool: List[Dict],
    num_samples: int = 5,
    num_shots: int = 3,
) -> tuple[str, int]:
    """
    Evaluate using self-consistency: sample multiple CoT paths and majority vote.
    Returns: (final_answer, agreement_count)
    The agreement count indicates how many of the num_samples paths agreed
    on the returned final_answer.
    """
    answers = []
    prompt = create_cot_prompt(train_pool, question, num_shots)

    # Simulate multiple independent samples with different random seeds
    for i in range(num_samples):
        random.seed(42 + i)  # Different seed for each sample
        response = simulate_model_response(question, ground_truth)
        answer = normalize_answer(extract_answer_from_text(response))
        if answer:
            answers.append(answer)

    # Majority vote: return the most common answer and its count
    if not answers:
        return "", 0

    answer_counts = Counter(answers)
    final_answer, count = answer_counts.most_common(1)[0]

    return final_answer, count


# Test self-consistency on a sample
test_item = test_list[1]
sc_answer, agreement = self_consistency_evaluate(
    test_item["question"], test_item["answer"], train_pool, num_samples=5
)

true_answer = normalize_answer(extract_answer_from_text(test_item["answer"]))