evaluate.py

"""
evaluate.py — Run PaderBot on the eval set and compute RAGAS-style metrics.

Output:
  - eval_predictions.json  — every question with its answer, sources, contexts
  - eval_scores.json       — per-question metric scores
  - eval_summary.csv       — aggregate table for the README

Three metrics:
  - Faithfulness:      fraction of claims supported by retrieved context
  - Answer relevance:  cosine sim between question and reverse-implied questions
  - Context precision: precision-at-rank averaged over relevant chunks

CACHING: predictions and scores are saved after every question. If the
script crashes (rate limits, etc.), just re-run — it picks up where it left off
stopped.

RATE LIMITING: uses the 8B model for evaluation calls + automatic retry
with backoff. Save the 70B budget for PaderBot's actual generation.

Usage:
    export GROQ_API_KEY=gsk_your_key
    python evaluate.py                  # full eval
    python evaluate.py --metrics-only   # use cached predictions, only re-score
    python evaluate.py --reset          # wipe cache and start fresh
"""

import argparse
import json
import os
import re
import sys
import time
from pathlib import Path

import numpy as np
from groq import Groq, RateLimitError
from sentence_transformers import SentenceTransformer

from paderbot import PaderBot
from eval_set import EVAL_SET

# ============================================================
# Config
# ============================================================
PREDICTIONS_PATH = Path("eval_predictions.json")
SCORES_PATH = Path("eval_scores.json")
SUMMARY_PATH = Path("eval_summary.csv")
EVAL_MODEL = "llama-3.1-8b-instant"
MAX_CONTEXT_CHARS = 3000   # truncate context per call to save tokens


# ============================================================
# Rate-limit-tolerant chat
# ============================================================
def make_chat(groq_client):
    def chat(messages, temperature=0, max_retries=5):
        time.sleep(0.4)
        for attempt in range(max_retries):
            try:
                resp = groq_client.chat.completions.create(
                    model=EVAL_MODEL, messages=messages, temperature=temperature, max_tokens=400,
                )
                return resp.choices[0].message.content
            except RateLimitError as e:
                wait = 2 ** attempt
                m = re.search(r"try again in ([\d.]+)s", str(e))
                if m:
                    wait = float(m.group(1)) + 0.5
                print(f"    [rate limit, waiting {wait:.1f}s]")
                time.sleep(wait)
        raise RuntimeError("Max retries exceeded — reduce eval set or wait longer.")
    return chat


# ============================================================
# Step 1: generate predictions
# ============================================================
def generate_predictions(bot, eval_set, cache):
    """Run PaderBot on every question, cache results."""
    for q in eval_set:
        if q["id"] in cache:
            print(f"  [cached] {q['id']}: {q['question'][:60]}")
            continue
        print(f"  Asking: {q['id']}: {q['question'][:60]}")
        try:
            result = bot.query(q["question"])
        except Exception as e:
            print(f"    ERROR: {e}")
            continue
        cache[q["id"]] = {
            "question": q["question"],
            "ground_truth": q["ground_truth"],
            "difficulty": q["difficulty"],
            "language": q["language"],
            "answer": result["answer"],
            "refused": result["refused"],
            "sources": result["sources"],
            "contexts": [{"text": c["text"], "url": c.get("url", "")} for c in result["contexts"]],
        }
        save_json(PREDICTIONS_PATH, cache)
    return cache


# ============================================================
# Step 2: metric implementations
# ============================================================
def split_into_claims(answer, chat):
    if "don't have enough information" in answer.lower() or "nicht genug" in answer.lower():
        return ["REFUSAL"]
    prompt = (
        "Break the following ANSWER into atomic factual claims (each a single verifiable statement). "
        "Output ONLY the claims, one per line, no numbering or commentary.\n\n"
        f"ANSWER: {answer}"
    )
    out = chat([{"role": "user", "content": prompt}])
    return [line.strip("- •*").strip() for line in out.splitlines() if line.strip()]


def is_supported(claim, context, chat):
    if len(context) > MAX_CONTEXT_CHARS:
        context = context[:MAX_CONTEXT_CHARS]
    prompt = (
        "Task: decide whether the CLAIM is supported by information in the CONTEXT.\n\n"
        "Both CONTEXT and CLAIM may be written in English or German. "
        "If the CONTEXT contains the same information as the CLAIM, even if paraphrased "
        "or in a different language, answer YES. Otherwise answer NO.\n\n"
        "Examples:\n"
        "CONTEXT: BAföG ist eine staatliche Förderung.\n"
        "CLAIM: BAföG is German state financial aid.\n"
        "Answer: YES\n\n"
        "CONTEXT: The library is open 24/7.\n"
        "CLAIM: Die Bibliothek hat begrenzte Öffnungszeiten.\n"
        "Answer: NO\n\n"
        f"CONTEXT:\n{context}\n\nCLAIM: {claim}\n\nAnswer (YES or NO only):"
        f"CONTEXT:\n{context}\n\nCLAIM: {claim}"
    )
    out = chat([{"role": "user", "content": prompt}]).strip().upper()
    return out.startswith("YES")


def faithfulness_score(pred, chat):
    claims = split_into_claims(pred["answer"], chat)
    if not claims or claims == ["REFUSAL"]:
        return None     # N/A for refusals
    context = "\n\n".join(c["text"] for c in pred["contexts"])
    if not context.strip():
        return None
    supported = sum(is_supported(c, context, chat) for c in claims)
    return supported / len(claims)


def generate_questions_from(answer, n, chat):
    prompt = (
        f"Below is an ANSWER. Generate {n} different questions that this answer would be a good response to. "
        "Output ONLY the questions, one per line, no numbering.\n\n"
        f"ANSWER: {answer}"
    )
    out = chat([{"role": "user", "content": prompt}])
    return [line.strip("- •*").strip() for line in out.splitlines() if line.strip()][:n]


def answer_relevance_score(pred, embedder, chat):
    answer = pred["answer"]
    if "don't have enough information" in answer.lower() or "nicht genug" in answer.lower():
        return None     # N/A for refusals
    generated = generate_questions_from(answer, 3, chat)
    if not generated:
        return 0.0
    q_emb = embedder.encode(pred["question"], normalize_embeddings=True)
    gen_embs = embedder.encode(generated, normalize_embeddings=True)
    sims = gen_embs @ q_emb
    return float(np.mean(sims))


def context_relevant(question, chunk_text, chat):
    chunk_preview = chunk_text[:MAX_CONTEXT_CHARS]
    prompt = (
        "Task: decide whether the CHUNK contains information that helps answer the QUESTION.\n\n"
        "Both QUESTION and CHUNK may be in English or German. "
        "If the CHUNK contains any information relevant to the QUESTION — even partially, "
        "or in a different language — answer YES. "
        "Only answer NO if the CHUNK is completely unrelated to the QUESTION.\n\n"
        "Examples:\n"
        "QUESTION: How do I apply for student housing?\n"
        "CHUNK: Über den Wohnplatzantrag kannst Du Dich direkt auf ein Zimmer bewerben.\n"
        "Answer: YES\n\n"
        "QUESTION: Was kostet die Miete?\n"
        "CHUNK: The Computer Science Master program covers algorithms and software engineering.\n"
        "Answer: NO\n\n"
        f"QUESTION: {question}\n\nCHUNK: {chunk_preview}\n\nAnswer (YES or NO only):"
    )
    out = chat([{"role": "user", "content": prompt}]).strip().upper()
    return out.startswith("YES")


def context_precision_score(pred, chat):
    if not pred["contexts"]:
        return 0.0
    relevances = [context_relevant(pred["question"], c["text"], chat) for c in pred["contexts"]]
    if not any(relevances):
        return 0.0
    precisions = []
    for k, is_rel in enumerate(relevances, start=1):
        if is_rel:
            precisions.append(sum(relevances[:k]) / k)
    return float(np.mean(precisions)) if precisions else 0.0


# ============================================================
# Step 3: score all predictions
# ============================================================
def score_all(predictions, embedder, chat, scores_cache):
    for pid, pred in predictions.items():
        if pid in scores_cache and all(
            k in scores_cache[pid] for k in ("faithfulness", "answer_relevance", "context_precision")
        ):
            print(f"  [cached] {pid}")
            continue
        print(f"  Scoring {pid}: {pred['question'][:60]}")
        try:
            faith = faithfulness_score(pred, chat)
            ar = answer_relevance_score(pred, embedder, chat)
            cp = context_precision_score(pred, chat)
        except Exception as e:
            print(f"    ERROR: {e}")
            continue
        scores_cache[pid] = {
            "id": pid,
            "difficulty": pred["difficulty"],
            "language": pred["language"],
            "refused": pred["refused"],
            "faithfulness": faith,
            "answer_relevance": ar,
            "context_precision": cp,
        }
        save_json(SCORES_PATH, scores_cache)
        # Friendly progress display
        f_s = f"{faith:.2f}" if faith is not None else "N/A"
        ar_s = f"{ar:.2f}" if ar is not None else "N/A"
        print(f"    faith={f_s}  rel={ar_s}  ctx_p={cp:.2f}")
    return scores_cache


# ============================================================
# Step 4: aggregate + write summary
# ============================================================
def summarize(scores):
    """Compute aggregate metrics overall + by language + by difficulty."""
    def mean(vals):
        clean = [v for v in vals if v is not None]
        return sum(clean) / len(clean) if clean else None

    def group_stats(items, label):
        return {
            "group": label,
            "n": len(items),
            "faithfulness": mean([s["faithfulness"] for s in items]),
            "answer_relevance": mean([s["answer_relevance"] for s in items]),
            "context_precision": mean([s["context_precision"] for s in items]),
            "refusal_rate": sum(1 for s in items if s["refused"]) / len(items) if items else 0,
        }

    all_scores = list(scores.values())
    rows = [group_stats(all_scores, "ALL")]
    rows.append(group_stats([s for s in all_scores if s["language"] == "en"], "EN"))
    rows.append(group_stats([s for s in all_scores if s["language"] == "de"], "DE"))
    rows.append(group_stats([s for s in all_scores if s["difficulty"] == "easy"], "easy"))
    rows.append(group_stats([s for s in all_scores if s["difficulty"] == "multi_hop"], "multi_hop"))
    rows.append(group_stats([s for s in all_scores if s["difficulty"] == "refusal"], "refusal"))
    return rows


def write_summary_csv(rows):
    import csv
    fields = ["group", "n", "faithfulness", "answer_relevance", "context_precision", "refusal_rate"]
    with SUMMARY_PATH.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for r in rows:
            row_out = {}
            for k in fields:
                v = r.get(k)
                row_out[k] = f"{v:.3f}" if isinstance(v, float) else v
            w.writerow(row_out)


def print_summary(rows):
    print("\n" + "=" * 70)
    print(f"{'Group':<12} {'N':>3}  {'Faith':>7} {'AnsRel':>7} {'CtxPrec':>8} {'Refusal':>8}")
    print("-" * 70)
    for r in rows:
        f = f"{r['faithfulness']:.2f}" if r["faithfulness"] is not None else "  N/A"
        ar = f"{r['answer_relevance']:.2f}" if r["answer_relevance"] is not None else "  N/A"
        cp = f"{r['context_precision']:.2f}" if r["context_precision"] is not None else "  N/A"
        ref = f"{r['refusal_rate']:.0%}"
        print(f"{r['group']:<12} {r['n']:>3}  {f:>7} {ar:>7} {cp:>8} {ref:>8}")
    print("=" * 70)


# ============================================================
# Cache helpers
# ============================================================
def load_json(path):
    if path.exists():
        with path.open() as f:
            return json.load(f)
    return {}


def save_json(path, data):
    with path.open("w") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


# ============================================================
# Main
# ============================================================
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--metrics-only", action="store_true", help="Use cached predictions; only re-score")
    parser.add_argument("--reset", action="store_true", help="Wipe caches and start fresh")
    args = parser.parse_args()

    if args.reset:
        for p in [PREDICTIONS_PATH, SCORES_PATH, SUMMARY_PATH]:
            if p.exists():
                p.unlink()
        print("Caches cleared.")

    if not os.environ.get("GROQ_API_KEY"):
        sys.exit("ERROR: GROQ_API_KEY not set in environment.")

    # Check for FILL IN placeholders in eval set
    fillins = [q["id"] for q in EVAL_SET if "[FILL IN" in q["ground_truth"]]
    if fillins:
        print(f"  WARNING: {len(fillins)} questions still have [FILL IN] ground_truth placeholders.")
        print(f"   IDs: {fillins[:5]}{'...' if len(fillins) > 5 else ''}")
        print(f"   Faithfulness metrics will still work (they use retrieved context, not ground_truth),")
        print(f"   but answer-correctness comparisons will be unreliable.")
        print(f"   Edit eval_set.py to fill these in.\n")
        if input("Continue anyway? [y/N] ").lower() != "y":
            sys.exit(0)

    # Set up clients
    groq = Groq()
    chat = make_chat(groq)
    embedder = SentenceTransformer("intfloat/multilingual-e5-base")

    # --- Step 1: predictions ---
    predictions = load_json(PREDICTIONS_PATH)
    if not args.metrics_only:
        print("\n=== Generating predictions ===")
        bot = PaderBot()
        predictions = generate_predictions(bot, EVAL_SET, predictions)

    # --- Step 2: scores ---
    print("\n=== Scoring with 3 metrics ===")
    scores = load_json(SCORES_PATH)
    scores = score_all(predictions, embedder, chat, scores)

    # --- Step 3: summarize ---
    rows = summarize(scores)
    write_summary_csv(rows)
    print_summary(rows)
    print(f"\n Wrote {PREDICTIONS_PATH}, {SCORES_PATH}, {SUMMARY_PATH}")


if __name__ == "__main__":
    main()