eval.py

#!/usr/bin/env python3
"""
Evaluation script for mini-grpo.

Evaluates a model on GSM8K test set or arithmetic problems and reports accuracy.

Usage:
    # Evaluate base GPT-2
    python eval.py --model gpt2 --dataset arithmetic --max_samples 50

    # Evaluate a GRPO-trained checkpoint
    python eval.py --model checkpoints/final --dataset arithmetic --max_samples 50

    # Compare base model vs trained model
    python eval.py --model gpt2 --trained_model checkpoints/final --dataset gsm8k --max_samples 100
"""

import argparse
import json
import logging
import os
from typing import Optional

import torch
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer

from mini_grpo.data import load_arithmetic, load_gsm8k
from mini_grpo.reward import extract_numeric_answer, math_correctness_reward

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("mini-grpo-eval")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate a model on math reasoning")

    parser.add_argument(
        "--model", type=str, default="gpt2",
        help="Base model name or path to evaluate.",
    )
    parser.add_argument(
        "--trained_model", type=str, default=None,
        help="Path to GRPO-trained checkpoint (for comparison).",
    )
    parser.add_argument(
        "--dataset", type=str, default="arithmetic",
        choices=["gsm8k", "arithmetic"],
        help="Dataset to evaluate on.",
    )
    parser.add_argument(
        "--max_samples", type=int, default=100,
        help="Number of evaluation samples.",
    )
    parser.add_argument(
        "--max_new_tokens", type=int, default=256,
        help="Max tokens to generate.",
    )
    parser.add_argument(
        "--temperature", type=float, default=0.1,
        help="Sampling temperature (low for evaluation).",
    )
    parser.add_argument(
        "--num_samples", type=int, default=1,
        help="Number of samples per problem (majority voting if > 1).",
    )
    parser.add_argument(
        "--device", type=str, default=None,
        help="Device (auto-detected if not set).",
    )
    parser.add_argument(
        "--output", type=str, default=None,
        help="Path to save evaluation results as JSON.",
    )
    parser.add_argument(
        "--dtype", type=str, default="float32",
        choices=["float32", "float16", "bfloat16"],
    )
    parser.add_argument(
        "--verbose", action="store_true", default=False,
        help="Print each problem and response.",
    )

    return parser.parse_args()


def detect_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def get_dtype(s: str) -> torch.dtype:
    return {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[s]


def load_model_for_eval(
    model_path: str,
    device: str,
    dtype: torch.dtype,
    base_model: Optional[str] = None,
):
    """Load a model for evaluation, handling both base models and LoRA checkpoints."""
    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        trust_remote_code=True,
        padding_side="left",
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    # Check if this is a PEFT checkpoint
    adapter_config_path = os.path.join(model_path, "adapter_config.json")
    if os.path.exists(adapter_config_path):
        from peft import PeftModel

        # Load adapter config to find base model
        with open(adapter_config_path, "r") as f:
            adapter_cfg = json.load(f)

        base_model_name = base_model or adapter_cfg.get("base_model_name_or_path", "gpt2")
        logger.info(f"Loading base model: {base_model_name}")

        base = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            torch_dtype=dtype,
            device_map=device,
            trust_remote_code=True,
        )
        model = PeftModel.from_pretrained(base, model_path)
        model.eval()
        # Re-load tokenizer from base if the checkpoint doesn't have one
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                model_path, trust_remote_code=True, padding_side="left"
            )
        except Exception:
            tokenizer = AutoTokenizer.from_pretrained(
                base_model_name, trust_remote_code=True, padding_side="left"
            )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=dtype,
            device_map=device,
            trust_remote_code=True,
        )
        model.eval()

    return model, tokenizer


@torch.no_grad()
def evaluate_model(
    model,
    tokenizer,
    dataset,
    max_new_tokens: int = 256,
    temperature: float = 0.1,
    num_samples: int = 1,
    device: str = "cuda",
    verbose: bool = False,
) -> dict:
    """
    Evaluate a model on a math reasoning dataset.

    Args:
        model: The language model.
        tokenizer: Tokenizer.
        dataset: MathReasoningDataset.
        max_new_tokens: Max generation length.
        temperature: Sampling temperature.
        num_samples: Samples per problem for majority voting.
        device: Device string.
        verbose: Print individual results.

    Returns:
        Dict with accuracy metrics and per-sample details.
    """
    model.eval()
    correct = 0
    total = 0
    details = []

    for idx in tqdm(range(len(dataset)), desc="Evaluating"):
        item = dataset[idx]
        prompt = item["prompt"]
        ground_truth = item["answer"]

        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=1024,
            padding=True,
        ).to(device)
        prompt_len = inputs["input_ids"].shape[1]

        # Generate potentially multiple samples
        sample_answers = []
        sample_texts = []
        for _ in range(num_samples):
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=max(temperature, 0.01),
                top_p=0.95,
                do_sample=temperature > 0.01,
                pad_token_id=tokenizer.pad_token_id,
            )
            generated = outputs[0, prompt_len:]
            text = tokenizer.decode(generated, skip_special_tokens=True)
            sample_texts.append(text)

            pred = extract_numeric_answer(text)
            sample_answers.append(pred)

        # Majority voting (or single answer)
        if num_samples == 1:
            is_correct = math_correctness_reward(sample_texts[0], ground_truth) > 0.5
            chosen_text = sample_texts[0]
        else:
            # Count valid numeric answers
            from collections import Counter
            valid = [a for a in sample_answers if a is not None]
            if valid:
                # Round for comparison
                rounded = [round(a, 2) for a in valid]
                most_common = Counter(rounded).most_common(1)[0][0]
                is_correct = math_correctness_reward(
                    f"<answer>{most_common}</answer>", ground_truth
                ) > 0.5
            else:
                is_correct = False
            chosen_text = sample_texts[0]

        if is_correct:
            correct += 1
        total += 1

        detail = {
            "index": idx,
            "question": item["question"],
            "ground_truth": ground_truth,
            "response": chosen_text[:500],  # Truncate for logging
            "predicted": extract_numeric_answer(chosen_text),
            "correct": is_correct,
        }
        details.append(detail)

        if verbose:
            status = "CORRECT" if is_correct else "WRONG"
            logger.info(
                f"[{status}] Q: {item['question'][:80]}... "
                f"GT: {ground_truth}, Pred: {detail['predicted']}"
            )

    accuracy = correct / total if total > 0 else 0.0
    return {
        "accuracy": accuracy,
        "correct": correct,
        "total": total,
        "details": details,
    }


def main():
    args = parse_args()

    if args.device is None:
        args.device = detect_device()
    dtype = get_dtype(args.dtype)

    # Load dataset
    logger.info(f"Loading evaluation dataset: {args.dataset}")
    if args.dataset == "gsm8k":
        dataset = load_gsm8k(split="test", max_samples=args.max_samples)
    else:
        dataset = load_arithmetic(
            n_samples=args.max_samples, seed=12345  # Different seed from training
        )
    logger.info(f"Evaluation samples: {len(dataset)}")

    # Evaluate base model
    logger.info(f"Loading base model: {args.model}")
    model, tokenizer = load_model_for_eval(args.model, args.device, dtype)

    logger.info("Evaluating base model...")
    base_results = evaluate_model(
        model, tokenizer, dataset,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        num_samples=args.num_samples,
        device=args.device,
        verbose=args.verbose,
    )
    logger.info(
        f"Base model accuracy: {base_results['accuracy']:.1%} "
        f"({base_results['correct']}/{base_results['total']})"
    )

    # Evaluate trained model if provided
    trained_results = None
    if args.trained_model is not None:
        logger.info(f"Loading trained model: {args.trained_model}")
        trained_model, trained_tokenizer = load_model_for_eval(
            args.trained_model, args.device, dtype, base_model=args.model
        )

        logger.info("Evaluating trained model...")
        trained_results = evaluate_model(
            trained_model, trained_tokenizer, dataset,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            num_samples=args.num_samples,
            device=args.device,
            verbose=args.verbose,
        )
        logger.info(
            f"Trained model accuracy: {trained_results['accuracy']:.1%} "
            f"({trained_results['correct']}/{trained_results['total']})"
        )

    # Print summary
    print("\n" + "=" * 60)
    print("EVALUATION RESULTS")
    print("=" * 60)
    print(f"Dataset        : {args.dataset} ({len(dataset)} samples)")
    print(f"Base model     : {args.model}")
    print(f"Base accuracy  : {base_results['accuracy']:.1%} "
          f"({base_results['correct']}/{base_results['total']})")

    if trained_results is not None:
        print(f"Trained model  : {args.trained_model}")
        print(f"Trained acc.   : {trained_results['accuracy']:.1%} "
              f"({trained_results['correct']}/{trained_results['total']})")
        delta = trained_results["accuracy"] - base_results["accuracy"]
        direction = "+" if delta >= 0 else ""
        print(f"Improvement    : {direction}{delta:.1%}")
    print("=" * 60)

    # Save results
    if args.output is not None:
        output_data = {
            "config": vars(args),
            "base_model": {
                "accuracy": base_results["accuracy"],
                "correct": base_results["correct"],
                "total": base_results["total"],
            },
        }
        if trained_results is not None:
            output_data["trained_model"] = {
                "accuracy": trained_results["accuracy"],
                "correct": trained_results["correct"],
                "total": trained_results["total"],
            }
        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
        with open(args.output, "w") as f:
            json.dump(output_data, f, indent=2)
        logger.info(f"Results saved to {args.output}")


if __name__ == "__main__":
    main()