eval_benchmark.py

"""
Genesis-2.0 RLHF — Evaluation Benchmark

Evaluates a model's tool-use capabilities on a hold-out set.
Measures: tool accuracy, completion rate, efficiency, format compliance.

Usage:
    python3 eval_benchmark.py --model <path> [--data <jsonl>] [--output <json>]
"""

import json
import os
import sys
import re
from typing import Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rewards import (
    extract_tool_calls,
    extract_assistant_messages,
    combined_reward,
    reward_debug,
)


def evaluate_trajectory(trajectory_text: str) -> dict:
    """
    Evaluate a single model trajectory on all metrics.
    Returns a dict of scores and metadata.
    """
    tool_calls = extract_tool_calls(trajectory_text)
    assistant_msgs = extract_assistant_messages(trajectory_text)
    reward_scores = reward_debug(trajectory_text)

    # Tool accuracy: how many tool calls have valid names + args
    valid_calls = 0
    for tc in tool_calls:
        name_ok = isinstance(tc.get("name"), str) and bool(tc["name"].strip())
        args_ok = "arguments" in tc and isinstance(tc["arguments"], dict)
        if name_ok and args_ok:
            valid_calls += 1

    tool_accuracy = valid_calls / max(len(tool_calls), 1)

    # Completion: is there a substantive answer?
    if assistant_msgs:
        last_msg = assistant_msgs[-1]
        has_text_answer = len(re.sub(r'<[^>]+>', '', last_msg).strip()) > 30
    else:
        has_text_answer = False

    return {
        "tool_count": len(tool_calls),
        "valid_tool_count": valid_calls,
        "tool_accuracy": tool_accuracy,
        "has_answer": 1.0 if has_text_answer else 0.0,
        "reward_completion": reward_scores["completion"],
        "reward_validity": reward_scores["validity"],
        "reward_efficiency": reward_scores["efficiency"],
        "reward_format": reward_scores["format"],
        "reward_combined": reward_scores["combined"],
        "assistant_turns": len(assistant_msgs),
    }


def evaluate_dataset(
    data_path: str,
    model_generate_fn=None,
    output_path: Optional[str] = None,
    max_samples: int = 200,
) -> dict:
    """
    Evaluate a dataset of prompts against a model.

    If model_generate_fn is None, evaluates the existing trajectories
    (for offline eval of SFT data).

    If model_generate_fn is provided, it's called as:
        response = model_generate_fn(prompt)
    for each prompt.

    Returns aggregated metrics.
    """
    # Load data
    data = []
    with open(data_path) as f:
        for i, line in enumerate(f):
            if i >= max_samples:
                break
            line = line.strip()
            if line:
                data.append(json.loads(line))

    results = []
    for item in data:
        if model_generate_fn:
            prompt = item.get("prompt", item.get("text", ""))
            response = model_generate_fn(prompt)
            full_text = prompt + "\n" + response
        else:
            full_text = item.get("text", "")

        metrics = evaluate_trajectory(full_text)
        results.append(metrics)

    # Aggregate
    agg = {
        "num_samples": len(results),
        "avg_tool_accuracy": sum(r["tool_accuracy"] for r in results) / max(len(results), 1),
        "avg_has_answer": sum(r["has_answer"] for r in results) / max(len(results), 1),
        "avg_reward_combined": sum(r["reward_combined"] for r in results) / max(len(results), 1),
        "avg_reward_completion": sum(r["reward_completion"] for r in results) / max(len(results), 1),
        "avg_reward_validity": sum(r["reward_validity"] for r in results) / max(len(results), 1),
        "avg_reward_efficiency": sum(r["reward_efficiency"] for r in results) / max(len(results), 1),
        "avg_reward_format": sum(r["reward_format"] for r in results) / max(len(results), 1),
        "avg_tool_count": sum(r["tool_count"] for r in results) / max(len(results), 1),
        "total_tool_calls": sum(r["tool_count"] for r in results),
    }

    if output_path:
        with open(output_path, "w") as f:
            json.dump({"aggregate": agg, "per_sample": results}, f, indent=2)
        print(f"Saved evaluation to {output_path}")

    return agg


def build_holdout_set(
    input_dir: str,
    output_path: str,
    num_prompts: int = 150,
    sources: Optional[list[str]] = None,
) -> None:
    """
    Build a hold-out evaluation set from SFT data.
    Takes prompts only (no completions) for online LLM evaluation.
    """
    if sources is None:
        sources = ["train_sessions_00001.jsonl", "train_augmented_00001.jsonl"]

    prompts = []
    for src in sources:
        path = os.path.join(input_dir, src)
        if not os.path.exists(path):
            print(f"  WARNING: {path} not found, skipping")
            continue
        with open(path) as f:
            for line in f:
                line = line.strip()
                if line:
                    item = json.loads(line)
                    # Extract just the prompt (user + system, no assistant)
                    text = item["text"]
                    idx = text.find("<|im_start|>assistant")
                    if idx >= 0:
                        prompt = text[:idx].strip()
                    else:
                        prompt = text
                    prompts.append({
                        "prompt": prompt,
                        "source": item["metadata"].get("source", "unknown"),
                    })

    # Take evenly spaced samples
    step = max(1, len(prompts) // num_prompts)
    holdout = [prompts[i] for i in range(0, len(prompts), step)][:num_prompts]

    with open(output_path, "w") as f:
        for p in holdout:
            f.write(json.dumps(p) + "\n")

    print(f"Built hold-out set: {len(holdout)} prompts → {output_path}")


if __name__ == "__main__":
    project_dir = os.path.dirname(os.path.abspath(__file__))
    data_dir = "/Volumes/this_and_that/hermes-admin/improvements/hermes-agentic-dataset/data/train"

    print("=== Genesis-2.0: Evaluation Benchmark ===\n")

    # Build hold-out set
    holdout_path = os.path.join(project_dir, "eval_holdout.jsonl")
    build_holdout_set(data_dir, holdout_path, num_prompts=150)
    print()

    # Evaluate SFT data offline (baseline)
    for src in ["train_sessions_00001.jsonl", "train_augmented_00001.jsonl"]:
        path = os.path.join(data_dir, src)
        if os.path.exists(path):
            print(f"\nEvaluating {src}...")
            agg = evaluate_dataset(path, max_samples=50)
            print(f"  Tool accuracy: {agg['avg_tool_accuracy']:.3f}")
            print(f"  Completion rate: {agg['avg_has_answer']:.3f}")
            print(f"  Combined reward: {agg['avg_reward_combined']:.3f}")
            print(f"  Tool efficiency: {agg['avg_reward_efficiency']:.3f}")
            print(f"  Format compliance: {agg['avg_reward_format']:.3f}")

    print(f"\nHold-out set ready: {holdout_path}")
    print(f"Run evaluation with: python3 eval_benchmark.py --eval {holdout_path}")