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()