evaluate.py

#!/usr/bin/env python3
"""
STEP 7 — Evaluate style fidelity against held-out validation pairs.
"""

from __future__ import annotations

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

import regex
import requests

from utils import load_config, read_jsonl, summarize_persona_for_prompt


def extract_emojis(text: str) -> set[str]:
    return set(regex.findall(r"\p{Extended_Pictographic}", text))


def token_set(text: str) -> set[str]:
    return set(re.findall(r"[a-zA-Z']+", text.lower()))


def completion_from_generated(full: str, prompt: str) -> str:
    if full.startswith(prompt):
        return full[len(prompt) :].strip()
    return full.strip()


def generate_ollama(
    prompt: str,
    model_name: str,
    temperature: float,
    base_url: str,
) -> str:
    resp = requests.post(
        f"{base_url}/api/generate",
        json={
            "model": model_name,
            "prompt": prompt,
            "stream": False,
            "options": {"temperature": temperature, "num_predict": 256},
        },
        timeout=120,
    )
    resp.raise_for_status()
    return resp.json().get("response", "")


def score_pair(actual: str, generated: str) -> dict:
    actual_emojis = extract_emojis(actual)
    gen_emojis = extract_emojis(generated)
    actual_tokens = token_set(actual)
    gen_tokens = token_set(generated)

    emoji_overlap = len(actual_emojis & gen_emojis) / max(len(actual_emojis), 1)
    vocab_overlap = len(actual_tokens & gen_tokens) / max(len(actual_tokens), 1)
    length_ratio = len(generated) / max(len(actual), 1)

    return {
        "length_actual": len(actual),
        "length_generated": len(generated),
        "length_ratio": round(length_ratio, 3),
        "vocab_overlap": round(vocab_overlap, 3),
        "emoji_overlap": round(emoji_overlap, 3),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--samples", type=int, default=20, help="Number of side-by-side comparisons")
    parser.add_argument("--ollama-url", default="http://localhost:11434")
    parser.add_argument("--dry-run", action="store_true", help="Score metrics only on empty generations")
    args = parser.parse_args()

    cfg = load_config()
    output_dir = Path(cfg["output_dir"])
    val_pairs_path = output_dir / "val_pairs.jsonl"
    profile_path = output_dir / "persona_profile.json"

    if not val_pairs_path.exists():
        raise FileNotFoundError(f"Missing {val_pairs_path}. Run build_dataset.py first.")

    pairs = read_jsonl(val_pairs_path)
    random.seed(42)
    sample_pairs = random.sample(pairs, min(args.samples, len(pairs)))

    model_name = cfg.get("ollama_model_name", "persona-imitate")
    temperature = cfg.get("default_temperature", 0.7)

    all_scores: list[dict] = []
    comparisons: list[dict] = []

    print(f"Evaluating {len(sample_pairs)} validation pairs via Ollama model '{model_name}' ...\n")

    for i, pair in enumerate(sample_pairs, 1):
        prompt = pair["context"]
        actual = pair["actual_reply"]

        if args.dry_run:
            generated = ""
        else:
            try:
                generated = generate_ollama(
                    prompt, model_name, temperature, args.ollama_url
                )
            except requests.RequestException as exc:
                print(f"Ollama request failed: {exc}")
                print("Start Ollama and run: ollama create persona-imitate -f Modelfile")
                print("Or use --dry-run to inspect val pairs without generation.")
                return

        scores = score_pair(actual, generated)
        all_scores.append(scores)
        comparisons.append({"actual": actual, "generated": generated, **scores})

        print(f"--- Sample {i} ({pair.get('period', '?')}) ---")
        print(f"ACTUAL:     {actual[:300]}")
        print(f"GENERATED:  {generated[:300]}")
        print(
            f"  len_ratio={scores['length_ratio']:.2f}  "
            f"vocab_overlap={scores['vocab_overlap']:.2f}  "
            f"emoji_overlap={scores['emoji_overlap']:.2f}"
        )
        print()

    if all_scores:
        avg = {
            key: round(sum(s[key] for s in all_scores) / len(all_scores), 3)
            for key in ["length_ratio", "vocab_overlap", "emoji_overlap"]
        }
        print("Aggregate metrics (sanity check, not accuracy):")
        print(f"  Avg length ratio:    {avg['length_ratio']}")
        print(f"  Avg vocab overlap:   {avg['vocab_overlap']}")
        print(f"  Avg emoji overlap:   {avg['emoji_overlap']}")

    report_path = output_dir / "eval_report.json"
    with open(report_path, "w", encoding="utf-8") as f:
        json.dump({"aggregate": avg if all_scores else {}, "comparisons": comparisons}, f, indent=2, ensure_ascii=False)
    print(f"\nReport saved -> {report_path}")


if __name__ == "__main__":
    main()