run_benchmark.py

#!/usr/bin/env python3
"""
Benchmark runner: sends each benchmark entry to a local Ollama model and saves responses.

Usage:
  python run_benchmark.py --model <model_name> [--questions ...] [--think] [--host <url>]

--questions selection rules:
  omitted             → all 18 questions
  --questions 5       → only EX5
  --questions 3 7     → EX3 through EX7 inclusive
  --questions 1 3 5 7 → exactly EX1, EX3, EX5, EX7
"""

import argparse
import json
import os
import re
import sys
import urllib.error
import urllib.request


# ---------------------------------------------------------------------------
# Paths (all relative to this script's directory)
# ---------------------------------------------------------------------------

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

BENCHMARK_JSONL = os.path.join(SCRIPT_DIR, "Dataset", "benchmark.jsonl")
DATASET_DIR = os.path.join(SCRIPT_DIR, "Dataset")
RESPONSES_DIR = os.path.join(SCRIPT_DIR, "Responses")

SYSTEM_PROMPT_TEMPLATE = os.path.join(SCRIPT_DIR, "LLMUnifiedConfig", "SystemPrompt.txt")
REFERENCE_TABLE        = os.path.join(SCRIPT_DIR, "LLMUnifiedConfig", "VEXVRReferenceTable.txt")
ASSIGNMENTS_DIR        = os.path.join(SCRIPT_DIR, "LLMUnifiedConfig", "Assignments")


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

def parse_args():
    parser = argparse.ArgumentParser(description="Run benchmark entries through an Ollama model.")
    parser.add_argument("--model", required=True, help="Ollama model name (e.g. qwen2.5:32b)")
    parser.add_argument(
        "--questions",
        nargs="+",
        type=int,
        metavar="N",
        help="Question IDs to run: single int, two ints (range), or 3+ ints (explicit list)",
    )
    parser.add_argument("--think", action="store_true", help="Enable thinking tokens (think: true)")
    parser.add_argument("--host", default="http://localhost:11434", help="Ollama base URL")
    return parser.parse_args()


def resolve_question_ids(questions_arg):
    """Convert --questions arg to a list of 'EX{N}' strings."""
    if questions_arg is None:
        return None  # means all
    if len(questions_arg) == 1:
        return [f"EX{questions_arg[0]}"]
    if len(questions_arg) == 2:
        lo, hi = questions_arg
        return [f"EX{n}" for n in range(lo, hi + 1)]
    return [f"EX{n}" for n in questions_arg]


# ---------------------------------------------------------------------------
# System prompt assembly
# ---------------------------------------------------------------------------

_prompt_cache = {}

def find_assignment_file(assignment_name):
    """Return the path to the assignment file, matching case-insensitively."""
    try:
        entries = os.listdir(ASSIGNMENTS_DIR)
    except FileNotFoundError:
        print(f"ERROR: Assignments directory not found: {ASSIGNMENTS_DIR}", file=sys.stderr)
        sys.exit(1)
    target = assignment_name.lower()
    for filename in entries:
        if filename.lower() == f"{target}.txt":
            return os.path.join(ASSIGNMENTS_DIR, filename)
    available = [f[:-4] for f in entries if f.endswith(".txt")]
    print(
        f"ERROR: No assignment file found for '{assignment_name}' in {ASSIGNMENTS_DIR}\n"
        f"  Available: {', '.join(sorted(available)) or '(none)'}",
        file=sys.stderr,
    )
    sys.exit(1)

def build_system_prompt(assignment_name):
    if assignment_name in _prompt_cache:
        return _prompt_cache[assignment_name]
    assignment_path = find_assignment_file(assignment_name)
    try:
        template   = open(SYSTEM_PROMPT_TEMPLATE, encoding="utf-8").read()
        reference  = open(REFERENCE_TABLE,         encoding="utf-8").read()
        assignment = open(assignment_path,          encoding="utf-8").read()
    except FileNotFoundError as e:
        print(f"ERROR: Could not read LLMUnifiedConfig file: {e}", file=sys.stderr)
        sys.exit(1)
    prompt = template.replace("{VEX_VR_REFERENCE}", reference).replace("{ASSIGNMENT}", assignment)
    _prompt_cache[assignment_name] = prompt
    return prompt


# ---------------------------------------------------------------------------
# User message construction
# ---------------------------------------------------------------------------

def build_user_message(entry):
    question = entry["student_question"]
    paths = entry["code_xml_paths"]

    xml_contents = []
    for path in paths:
        full_path = os.path.join(DATASET_DIR, path)
        if not os.path.exists(full_path):
            raise FileNotFoundError(f"XML file not found: {full_path}")
        xml_contents.append(open(full_path, encoding="utf-8").read())

    if len(xml_contents) == 1:
        code_section = xml_contents[0]
    else:
        parts = []
        for i, content in enumerate(xml_contents, 1):
            parts.append(f"[Attempt {i}]\n{content}")
        code_section = "\n\n".join(parts)

    return f"{question}\n\n--- Code ---\n{code_section}"


# ---------------------------------------------------------------------------
# Ollama API call
# ---------------------------------------------------------------------------

def call_ollama(host, model, system_prompt, user_message, think):
    url = f"{host.rstrip('/')}/api/chat"
    body = {
        "model": model,
        "stream": False,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
    }
    if think:
        body["think"] = True

    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})

    try:
        with urllib.request.urlopen(req) as resp:
            raw = resp.read().decode("utf-8")
    except urllib.error.URLError as e:
        raise ConnectionError(f"Could not reach Ollama at {host}: {e.reason}") from e
    except urllib.error.HTTPError as e:
        body_text = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {body_text}") from e

    return json.loads(raw)


# ---------------------------------------------------------------------------
# Response file formatting
# ---------------------------------------------------------------------------

def format_response_file(elapsed_s, thinking, response):
    lines = [f"Time: {elapsed_s:.2f}s"]
    if thinking:
        lines += ["", "[THINKING]", thinking, "[/THINKING]"]
    lines += ["", "[RESPONSE]", response, "[/RESPONSE]"]
    return "\n".join(lines) + "\n"


# ---------------------------------------------------------------------------
# Trial directory helpers
# ---------------------------------------------------------------------------

def model_slug(model_name):
    return re.sub(r"[:/\\]", "_", model_name)


def next_trial_dir(model_dir):
    os.makedirs(model_dir, exist_ok=True)
    existing = [
        d for d in os.listdir(model_dir)
        if re.match(r"^Trial\d+$", d) and os.path.isdir(os.path.join(model_dir, d))
    ]
    if not existing:
        return os.path.join(model_dir, "Trial1"), 1
    nums = [int(re.search(r"\d+", d).group()) for d in existing]
    n = max(nums) + 1
    return os.path.join(model_dir, f"Trial{n}"), n


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()

    # Resolve question IDs
    selected_ids = resolve_question_ids(args.questions)

    # Load benchmark entries
    entries = []
    with open(BENCHMARK_JSONL, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                entries.append(json.loads(line))

    if selected_ids is not None:
        id_set = set(selected_ids)
        entries = [e for e in entries if e["id"] in id_set]
        missing = id_set - {e["id"] for e in entries}
        if missing:
            print(f"WARNING: IDs not found in benchmark: {sorted(missing)}", file=sys.stderr)

    if not entries:
        print("No matching entries found. Exiting.", file=sys.stderr)
        sys.exit(1)

    # Determine output directory
    slug = model_slug(args.model)
    model_dir = os.path.join(RESPONSES_DIR, slug)
    trial_dir, trial_num = next_trial_dir(model_dir)
    os.makedirs(trial_dir, exist_ok=True)

    print(f"Model: {args.model} | Trial: {trial_num} | Questions: {len(entries)}")
    print(f"Output: {trial_dir}")

    # Save system prompt snapshot for reproducibility
    prompt_snap = build_system_prompt(entries[0]["assignment"])
    with open(os.path.join(trial_dir, "system_prompt.txt"), "w", encoding="utf-8") as f:
        f.write(prompt_snap)

    # Check Ollama reachability before starting
    try:
        req = urllib.request.Request(f"{args.host.rstrip('/')}/api/tags")
        with urllib.request.urlopen(req):
            pass
    except Exception as e:
        print(f"\nERROR: Cannot reach Ollama at {args.host}\n  {e}", file=sys.stderr)
        sys.exit(1)

    # Run each entry
    total_time = 0.0
    saved = 0

    for entry in entries:
        eid = entry["id"]
        out_path = os.path.join(trial_dir, f"{eid}.txt")

        # Build system prompt for this entry's assignment
        system_prompt = build_system_prompt(entry["assignment"])

        # Build user message
        try:
            user_message = build_user_message(entry)
        except FileNotFoundError as e:
            print(f"  ✗ {eid} — SKIPPED: {e}")
            continue

        # Call Ollama
        try:
            result = call_ollama(args.host, args.model, system_prompt, user_message, args.think)
        except ConnectionError as e:
            print(f"\nERROR: {e}", file=sys.stderr)
            sys.exit(1)
        except RuntimeError as e:
            content = f"ERROR: {e}\n"
            with open(out_path, "w", encoding="utf-8") as f:
                f.write(content)
            print(f"  ✗ {eid} — API error (saved): {e}")
            continue

        # Parse result
        elapsed_s = result.get("total_duration", 0) / 1e9
        total_time += elapsed_s
        msg = result.get("message", {})
        response_text = msg.get("content", "")
        thinking_text = msg.get("thinking") or ""

        # Write response file immediately
        file_content = format_response_file(elapsed_s, thinking_text, response_text)
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(file_content)

        saved += 1
        think_note = " (with thinking)" if thinking_text else ""
        print(f"  ✓ {eid} ({elapsed_s:.2f}s){think_note}")

    print(f"\nDone. {saved}/{len(entries)} response(s) saved to {trial_dir}")
    print(f"Total time: {total_time:.1f}s")


if __name__ == "__main__":
    main()