benchmark.py

"""Benchmark AccountantGPT against CPA FAR questions.

Usage:
    # Benchmark fine-tuned Gemma 4 model
    python benchmark.py --model ./models/accountant-gpt-gemma4-31b-combined/final \
        --base-model google/gemma-4-31b-it --n 500

    # Benchmark Qwen2.5 adapter
    python benchmark.py --model ./models/accountant-gpt-combined/final \
        --base-model Qwen/Qwen2.5-7B-Instruct --n 500

    # Benchmark base model (no adapter) for comparison
    python benchmark.py --model Qwen/Qwen2.5-7B-Instruct --no-adapter --n 500
"""

import argparse
import json
import random
import re
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

SCRIPT_DIR = Path(__file__).resolve().parent
DATA_DIR = SCRIPT_DIR.parent / "data"


def load_cpa_questions(n: int = 500) -> list[dict]:
    path = DATA_DIR / "cpa-far" / "questions.jsonl"
    all_qs = []
    with open(path) as f:
        for line in f:
            if line.strip():
                all_qs.append(json.loads(line))
    random.seed(42)
    return random.sample(all_qs, min(n, len(all_qs)))


def format_cpa_prompt(q: dict) -> str:
    choices = "\n".join(f"{k}) {v}" for k, v in sorted(q["choices"].items()))
    return (
        f"Answer this CPA exam question. State ONLY the correct letter "
        f"(A, B, C, or D) on the first line, then explain briefly.\n\n"
        f"{q['question']}\n\n{choices}"
    )


def extract_answer_letter(response: str) -> str:
    response = response.strip()
    first_line = response.split("\n")[0].strip()
    match = re.search(r"\b([A-D])\b", first_line)
    if match:
        return match.group(1)
    match = re.search(r"(?:answer|correct)\s+(?:is\s+)?([A-D])\b", response, re.I)
    if match:
        return match.group(1)
    match = re.search(r"\b([A-D])\)", response)
    if match:
        return match.group(1)
    return ""


def load_model(model_path: str, base_model: str, is_adapter: bool = True):
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    if is_adapter:
        model = AutoModelForCausalLM.from_pretrained(
            base_model,
            quantization_config=bnb_config,
            device_map="auto",
            torch_dtype=torch.bfloat16,
        )
        model = PeftModel.from_pretrained(model, model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_path)
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            quantization_config=bnb_config,
            device_map="auto",
            torch_dtype=torch.bfloat16,
        )
        tokenizer = AutoTokenizer.from_pretrained(model_path)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer


def generate(model, tokenizer, prompt: str, system: str = "", max_new_tokens: int = 512) -> str:
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.1,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
        )

    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return response.strip()


def run_benchmark(model, tokenizer, n: int = 500):
    questions = load_cpa_questions(n)
    correct = 0
    total = 0
    results_by_topic = {}

    system = (
        "You are an expert CPA exam tutor. Answer multiple-choice accounting "
        "questions accurately and concisely."
    )

    print(f"\nBenchmarking on {len(questions)} CPA FAR questions...")
    for i, q in enumerate(questions):
        prompt = format_cpa_prompt(q)
        response = generate(model, tokenizer, prompt, system=system)
        predicted = extract_answer_letter(response)
        is_correct = predicted == q["correct_answer"]

        if is_correct:
            correct += 1
        total += 1

        topic = q.get("topic", "Unknown")
        if topic not in results_by_topic:
            results_by_topic[topic] = {"correct": 0, "total": 0}
        results_by_topic[topic]["total"] += 1
        if is_correct:
            results_by_topic[topic]["correct"] += 1

        if (i + 1) % 50 == 0:
            print(f"  [{i+1}/{len(questions)}] Running accuracy: {correct/total:.1%}")

    print(f"\n{'='*60}")
    print(f"CPA FAR Benchmark Results")
    print(f"{'='*60}")
    print(f"Overall: {correct}/{total} = {correct/total:.1%}")

    topic_scores = [
        (t, d["correct"] / d["total"], d["total"])
        for t, d in results_by_topic.items()
        if d["total"] >= 3
    ]
    topic_scores.sort(key=lambda x: x[1], reverse=True)

    if topic_scores:
        print("\nTop 5 topics:")
        for t, acc, n_q in topic_scores[:5]:
            print(f"  {acc:.0%} ({n_q} qs) - {t}")
        print("\nBottom 5 topics:")
        for t, acc, n_q in topic_scores[-5:]:
            print(f"  {acc:.0%} ({n_q} qs) - {t}")

    results = {
        "overall_accuracy": correct / total,
        "total_questions": total,
        "correct": correct,
        "by_topic": {t: {"accuracy": d["correct"] / d["total"], **d}
                     for t, d in results_by_topic.items()},
    }
    out_path = SCRIPT_DIR / "benchmark_results.json"
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {out_path}")
    return results


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True, help="Path to adapter or full model")
    parser.add_argument("--base-model", type=str, default="Qwen/Qwen2.5-7B-Instruct")
    parser.add_argument("--no-adapter", action="store_true")
    parser.add_argument("--n", type=int, default=500)
    args = parser.parse_args()

    print(f"Loading model: {args.model}")
    if not args.no_adapter:
        print(f"Base model: {args.base_model}")
    model, tokenizer = load_model(args.model, args.base_model, is_adapter=not args.no_adapter)

    run_benchmark(model, tokenizer, n=args.n)


if __name__ == "__main__":
    main()