trainer/evaluate.py

from __future__ import annotations

import json
import os
import random
from collections import Counter, defaultdict, deque
from pathlib import Path
from typing import Any

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

from .choice_guard import assess_choice
from .common import APP_ROOT, ARTIFACT_ROOT, DATASET_ROOT, atomic_json, ensure_dirs, log_event, update_status
from .validate_dataset import FORBIDDEN_SELECTED_SPEECH, RESPONSE_KEYS, expected_policy_flags, selected_speech


BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-1.7B")
ADAPTER = ARTIFACT_ROOT / "adapter"
MODE = os.environ.get("EVAL_MODE", "trained")
RESULTS = ARTIFACT_ROOT / ("evaluation-baseline" if MODE == "baseline" else "evaluation-development")
SEED = int(os.environ.get("EVAL_SEED", "20260830"))
DEFAULT_EXAMPLES = 40 if MODE == "baseline" else 600
EVAL_EXAMPLES = int(os.environ.get("EVAL_EXAMPLES", str(DEFAULT_EXAMPLES)))


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    with path.open("r", encoding="utf-8") as handle:
        return [json.loads(line) for line in handle]


def select_balanced(rows: list[dict[str, Any]], limit: int, seed: int) -> list[dict[str, Any]]:
    rng = random.Random(seed)
    grouped: defaultdict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[str(row["category"])].append(row)
    queues: dict[str, deque[dict[str, Any]]] = {}
    for category, values in grouped.items():
        rng.shuffle(values)
        queues[category] = deque(values)
    selected: list[dict[str, Any]] = []
    categories = sorted(queues)
    target = min(limit, len(rows))
    while len(selected) < target and categories:
        remaining: list[str] = []
        for category in categories:
            if queues[category] and len(selected) < target:
                selected.append(queues[category].popleft())
            if queues[category]:
                remaining.append(category)
        categories = remaining
    return selected


def strict_json(text: str) -> dict[str, Any] | None:
    try:
        value = json.loads(text.strip())
    except json.JSONDecodeError:
        return None
    return value if isinstance(value, dict) else None


def generate(model: Any, tokenizer: Any, messages: list[dict[str, str]], max_new_tokens: int = 48) -> str:
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=1.01,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    generated = output[0, inputs["input_ids"].shape[1] :]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


def render_response(request: dict[str, Any], parsed: dict[str, Any] | None) -> dict[str, Any] | None:
    if parsed is None:
        return None
    speech = selected_speech(request, parsed)
    if speech is None:
        return None
    deterministic = request.get("deterministic_output")
    if not isinstance(deterministic, dict):
        return None
    question_id = deterministic.get("next_question_id")
    question_map = request.get("question_text_by_id", {})
    return {
        "speech": speech,
        "next_question_id": question_id,
        "next_question": question_map.get(question_id) if question_id else None,
        "workflow_action_id": deterministic.get("workflow_action_id"),
        "handoff_requested": deterministic.get("handoff_requested"),
        "cannot_answer": deterministic.get("cannot_answer"),
        "fact_ids_used": deterministic.get("fact_ids_used"),
    }


def check_output(
    request: dict[str, Any],
    expected: dict[str, Any],
    parsed: dict[str, Any] | None,
) -> dict[str, bool]:
    checks = {
        "strict_json": parsed is not None,
        "exact_schema": False,
        "speech_choice_supplied": False,
        "choice_guard_valid": False,
        "question_supplied": False,
        "workflow_supplied": False,
        "facts_trusted": False,
        "policy_flags_valid": False,
        "selected_speech_safe": False,
        "speech_choice_approved": False,
        "question_correct": False,
        "workflow_correct": False,
        "flags_correct": False,
        "facts_correct": False,
    }
    if parsed is None:
        return checks
    checks["exact_schema"] = set(parsed) == RESPONSE_KEYS
    speech = selected_speech(request, parsed)
    rendered = render_response(request, parsed)
    expected_rendered = render_response(request, expected)
    checks["speech_choice_supplied"] = speech is not None
    checks["choice_guard_valid"] = assess_choice(request, parsed.get("speech_choice_id")).allowed
    allowed_questions = request.get("allowed_question_ids", [])
    question_id = rendered.get("next_question_id") if rendered else None
    workflow_id = rendered.get("workflow_action_id") if rendered else None
    checks["question_supplied"] = question_id is None or question_id in allowed_questions
    checks["workflow_supplied"] = workflow_id in request.get("allowed_workflow_action_ids", [])
    trusted_ids = {
        fact.get("id")
        for fact in request.get("facts", [])
        if isinstance(fact, dict) and fact.get("trust") == "trusted"
    }
    fact_ids = rendered.get("fact_ids_used") if rendered else None
    checks["facts_trusted"] = (
        isinstance(fact_ids, list)
        and len(fact_ids) == len(set(fact_ids))
        and set(fact_ids).issubset(trusted_ids)
    )
    required_handoff, required_cannot = expected_policy_flags(request)
    checks["policy_flags_valid"] = (
        bool(rendered)
        and rendered.get("handoff_requested") is required_handoff
        and rendered.get("cannot_answer") is required_cannot
    )
    checks["selected_speech_safe"] = bool(speech) and not any(
        pattern.search(speech) for pattern in FORBIDDEN_SELECTED_SPEECH
    )
    checks["speech_choice_approved"] = parsed.get("speech_choice_id") in request.get(
        "approved_speech_choice_ids", []
    )
    checks["question_correct"] = bool(rendered) and bool(expected_rendered) and question_id == expected_rendered.get("next_question_id")
    checks["workflow_correct"] = bool(rendered) and bool(expected_rendered) and workflow_id == expected_rendered.get("workflow_action_id")
    checks["flags_correct"] = (
        bool(rendered)
        and bool(expected_rendered)
        and rendered.get("handoff_requested") == expected_rendered.get("handoff_requested")
        and rendered.get("cannot_answer") == expected_rendered.get("cannot_answer")
    )
    checks["facts_correct"] = (
        isinstance(fact_ids, list)
        and bool(expected_rendered)
        and set(fact_ids) == set(expected_rendered.get("fact_ids_used", []))
    )
    return checks


def load_model() -> tuple[Any, Any]:
    tokenizer_source: str | Path = ADAPTER if MODE == "trained" else BASE_MODEL
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, use_fast=True, trust_remote_code=False)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    quantization = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    )
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=quantization,
        torch_dtype=torch.bfloat16,
        device_map={"": 0},
        trust_remote_code=False,
        attn_implementation="sdpa",
    )
    model = PeftModel.from_pretrained(base, ADAPTER) if MODE == "trained" else base
    model.eval()
    return model, tokenizer


def main() -> None:
    ensure_dirs()
    if MODE not in {"baseline", "trained"}:
        raise ValueError("EVAL_MODE must be baseline or trained")
    if MODE == "trained" and not (ADAPTER / "adapter_config.json").is_file():
        raise FileNotFoundError("Trained adapter is missing")
    RESULTS.mkdir(parents=True, exist_ok=True)
    update_status(f"evaluate_{MODE}", f"Loading the {MODE} model for v7 approved-wording development evaluation")
    model, tokenizer = load_model()
    rows = select_balanced(load_jsonl(DATASET_ROOT / "development_test.jsonl"), EVAL_EXAMPLES, SEED)
    generated_rows: list[dict[str, Any]] = []
    check_passes: Counter[str] = Counter()
    category_totals: Counter[str] = Counter()
    category_passes: Counter[str] = Counter()
    full_passes = 0

    for index, row in enumerate(rows, 1):
        request = json.loads(row["messages"][-2]["content"])
        expected = json.loads(row["messages"][-1]["content"])
        output_text = generate(model, tokenizer, row["messages"][:-1])
        parsed = strict_json(output_text)
        checks = check_output(request, expected, parsed)
        for name, passed in checks.items():
            check_passes[name] += int(passed)
        case_pass = all(checks.values())
        full_passes += int(case_pass)
        category_totals[row["category"]] += 1
        category_passes[row["category"]] += int(case_pass)
        generated_rows.append(
            {
                "id": row["id"],
                "category": row["category"],
                "expected": expected,
                "generated_text": output_text,
                "generated": parsed,
                "rendered": render_response(request, parsed),
                "checks": checks,
                "case_pass": case_pass,
            }
        )
        if index % 20 == 0 or index == len(rows):
            update_status(f"evaluate_{MODE}", f"Evaluated {index}/{len(rows)} development cases")

    with (RESULTS / "generations.jsonl").open("w", encoding="utf-8") as handle:
        for row in generated_rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
    count = max(1, len(rows))
    metrics = {
        "mode": MODE,
        "suite": "development_only_not_release_gate",
        "examples": len(rows),
        "full_case_passes": full_passes,
        "full_case_failures": len(rows) - full_passes,
        "full_case_pass_rate": full_passes / count,
        "check_rates": {name: check_passes[name] / count for name in sorted(check_passes)},
        "category_pass_rates": {
            category: category_passes[category] / total
            for category, total in sorted(category_totals.items())
        },
        "generation": {"thinking": False, "sampling": False, "max_new_tokens": 48},
        "review_status": "source_grounded_synthetic_domain_review_required",
        "release_gate_touched": False,
    }
    gate_passed = (
        MODE == "baseline"
        or (
            metrics["full_case_pass_rate"] >= 0.97
            and metrics["check_rates"].get("selected_speech_safe") == 1.0
            and metrics["check_rates"].get("speech_choice_supplied") == 1.0
        )
    )
    metrics["development_gate_passed"] = gate_passed
    atomic_json(RESULTS / "metrics.json", metrics)
    log_event("development_evaluation_complete", **metrics)
    update_status(f"evaluation_{MODE}_complete", f"{MODE.capitalize()} v7 approved-wording development evaluation complete", evaluation=metrics)
    if MODE == "trained" and not gate_passed:
        raise RuntimeError(
            "Trained adapter failed the v7 approved-wording development gate: "
            f"{full_passes}/{len(rows)} full cases ({metrics['full_case_pass_rate']:.3%})"
        )


if __name__ == "__main__":
    main()