eval.py

#!/usr/bin/env python3
"""
WASSCE Benchmark Evaluation Script

Evaluates open language models against the West African Senior School Certificate
Examination (WASSCE) using the Situated Evaluation Framework.

Usage:
    python eval.py --model deepseek-ai/DeepSeek-V3-0324 --config configs/crusoe.yaml
    python eval.py --model deepseek-ai/DeepSeek-V3-0324 --subject core_maths --config configs/crusoe.yaml
    python eval.py --all-models --config configs/crusoe.yaml
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import requests
import yaml

SUBJECTS = ["core_maths", "english", "integrated_science", "social_studies"]

SUBJECT_DISPLAY = {
    "core_maths": "Core Mathematics",
    "english": "English Language",
    "integrated_science": "Integrated Science",
    "social_studies": "Social Studies",
}

PROMPT_TEMPLATE = """The following is a multiple choice question from the West African Senior School Certificate Examination (WASSCE) in {subject}.

Question: {question}

A. {option_a}
B. {option_b}
C. {option_c}
D. {option_d}

Answer with just the letter of the correct option."""


def load_config(config_path: str) -> dict:
    with open(config_path) as f:
        return yaml.safe_load(f)


def load_questions(subject: str, data_dir: str = "data") -> list[dict]:
    path = Path(data_dir) / subject / "questions.json"
    if not path.exists():
        print(f"Warning: No questions found for {subject} at {path}")
        return []
    with open(path) as f:
        return json.load(f)


def build_prompt(question: dict) -> str:
    return PROMPT_TEMPLATE.format(
        subject=SUBJECT_DISPLAY.get(question["subject"], question["subject"]),
        question=question["question"],
        option_a=question["options"]["A"],
        option_b=question["options"]["B"],
        option_c=question["options"]["C"],
        option_d=question["options"]["D"],
    )


def extract_answer(response_text: str) -> str | None:
    text = response_text.strip().upper()
    for char in ["A", "B", "C", "D"]:
        if text == char or text.startswith(f"{char}.") or text.startswith(f"{char})"):
            return char
    for char in ["A", "B", "C", "D"]:
        if char in text:
            return char
    return None


def query_model(
    prompt: str, model_id: str, config: dict, max_retries: int = 3
) -> str:
    base_url = config["provider"]["base_url"].rstrip("/")
    api_key = config["provider"]["api_key"]

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": model_id,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 16,
        "temperature": 0.0,
    }

    for attempt in range(max_retries):
        try:
            response = requests.post(
                f"{base_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=60,
            )
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"]
        except requests.exceptions.RequestException as e:
            if attempt < max_retries - 1:
                wait = 2 ** (attempt + 1)
                print(f"  Retry {attempt + 1}/{max_retries} after {wait}s: {e}")
                time.sleep(wait)
            else:
                print(f"  Failed after {max_retries} attempts: {e}")
                return ""


def evaluate_subject(
    subject: str, model_id: str, config: dict, data_dir: str = "data"
) -> dict:
    questions = load_questions(subject, data_dir)
    if not questions:
        return {"subject": subject, "total": 0, "correct": 0, "accuracy": 0.0, "details": []}

    correct = 0
    details = []

    print(f"\n  {SUBJECT_DISPLAY.get(subject, subject)} ({len(questions)} questions)")

    for i, q in enumerate(questions):
        prompt = build_prompt(q)
        response = query_model(prompt, model_id, config)
        predicted = extract_answer(response)
        is_correct = predicted == q["correct_answer"]

        if is_correct:
            correct += 1

        details.append({
            "id": q.get("id", f"{subject}_{i+1}"),
            "question_number": q.get("question_number", i + 1),
            "correct_answer": q["correct_answer"],
            "predicted_answer": predicted,
            "raw_response": response,
            "is_correct": is_correct,
            "topic": q.get("topic", "unknown"),
        })

        status = "+" if is_correct else "x"
        print(f"    [{status}] Q{q.get('question_number', i+1)}: {predicted or '?'} (correct: {q['correct_answer']})")

    accuracy = correct / len(questions) if questions else 0.0
    print(f"  Result: {correct}/{len(questions)} ({accuracy:.1%})")

    return {
        "subject": subject,
        "total": len(questions),
        "correct": correct,
        "accuracy": round(accuracy, 4),
        "details": details,
    }


def evaluate_model(
    model_id: str,
    config: dict,
    subjects: list[str] | None = None,
    data_dir: str = "data",
    output_dir: str = "results",
) -> dict:
    subjects = subjects or SUBJECTS
    print(f"\nEvaluating: {model_id}")
    print(f"Subjects: {', '.join(subjects)}")
    print("-" * 60)

    results = {
        "model": model_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "framework": "situated_evaluation_framework",
        "subjects": {},
        "summary": {},
    }

    total_correct = 0
    total_questions = 0

    for subject in subjects:
        subject_result = evaluate_subject(subject, model_id, config, data_dir)
        results["subjects"][subject] = subject_result
        total_correct += subject_result["correct"]
        total_questions += subject_result["total"]

    overall_accuracy = total_correct / total_questions if total_questions else 0.0

    results["summary"] = {
        "total_questions": total_questions,
        "total_correct": total_correct,
        "overall_accuracy": round(overall_accuracy, 4),
        "subject_accuracies": {
            s: results["subjects"][s]["accuracy"] for s in subjects
        },
    }

    # Print summary
    print("\n" + "=" * 60)
    print(f"RESULTS: {model_id}")
    print("=" * 60)
    for subject in subjects:
        r = results["subjects"][subject]
        bar = "#" * int(r["accuracy"] * 20) + "-" * (20 - int(r["accuracy"] * 20))
        print(f"  {SUBJECT_DISPLAY.get(subject, subject):25s} [{bar}] {r['accuracy']:.1%} ({r['correct']}/{r['total']})")
    print(f"  {'OVERALL':25s} {'':22s} {overall_accuracy:.1%} ({total_correct}/{total_questions})")
    print("=" * 60)

    # Save results
    os.makedirs(output_dir, exist_ok=True)
    model_slug = model_id.replace("/", "__")
    output_path = Path(output_dir) / f"{model_slug}.json"
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to: {output_path}")

    return results


def main():
    parser = argparse.ArgumentParser(description="WASSCE Benchmark Evaluation")
    parser.add_argument("--model", type=str, help="Model ID to evaluate")
    parser.add_argument("--all-models", action="store_true", help="Evaluate all models in config")
    parser.add_argument("--subject", type=str, choices=SUBJECTS, help="Evaluate a single subject")
    parser.add_argument("--config", type=str, required=True, help="Path to config YAML")
    parser.add_argument("--data-dir", type=str, default="data", help="Path to data directory")
    parser.add_argument("--output-dir", type=str, default="results", help="Path to output directory")
    args = parser.parse_args()

    if not args.model and not args.all_models:
        parser.error("Specify --model MODEL_ID or --all-models")

    config = load_config(args.config)
    subjects = [args.subject] if args.subject else SUBJECTS

    if args.all_models:
        models = config.get("models", {}).keys()
        if not models:
            print("Error: No models defined in config")
            sys.exit(1)
        for model_id in models:
            evaluate_model(model_id, config, subjects, args.data_dir, args.output_dir)
    else:
        evaluate_model(args.model, config, subjects, args.data_dir, args.output_dir)


if __name__ == "__main__":
    main()