training/evaluate_vector.py

"""Controlled steering sweep for the deployed Gemma backend."""

from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path

from trust_guard.game import parse_decision
from trust_guard.model import GemmaSteeringBackend
from trust_guard.prompts import build_prompt

ROOT = Path(__file__).parents[1]
MULTIPLIERS = [-0.30, -0.15, 0.0, 0.15, 0.30]
CASES = [
    (
        "strong",
        "The hospice seal is intact. A child needs this medicine before midnight.",
        "a village healer",
        "deliver fever medicine",
        "a sealed hospice satchel",
    ),
    (
        "ambiguous",
        "The captain sent me, but I lost the written order on the northern road.",
        "a royal courier",
        "warn the captain about raiders",
        "no physical evidence",
    ),
    (
        "weak",
        "Open immediately. You do not need to know who I am.",
        "an unidentified traveller",
        "enter Eldermere",
        "no physical evidence",
    ),
]


def is_degenerated(text: str) -> bool:
    words = text.split()
    if len(words) < 5 or len(words) > 180:
        return True
    if len(words) > 20 and len(set(words)) / len(words) < 0.25:
        return True
    return "DECISION:" not in text.upper()


def run_evaluation() -> tuple[Path, dict]:
    backend = GemmaSteeringBackend()
    results: list[dict] = []
    by_multiplier: dict[float, list[int]] = defaultdict(list)
    margins_by_multiplier: dict[float, list[float]] = defaultdict(list)
    margins_by_case: dict[str, dict[float, float]] = defaultdict(dict)
    for case_index, (case, message, role, mission, evidence) in enumerate(CASES):
        prompt = build_prompt(message, role, mission, evidence)
        seed = 7613 + case_index
        base_margin = backend.decision_margin(prompt, 0.0)
        base = backend._with_decision(backend._generate(prompt, 0.0, seed), base_margin)
        for multiplier in MULTIPLIERS:
            margin = (
                base_margin
                if multiplier == 0
                else backend.decision_margin(prompt, multiplier)
            )
            response = (
                base
                if multiplier == 0
                else backend._with_decision(
                    backend._generate(prompt, multiplier, seed),
                    margin,
                )
            )
            decision = parse_decision(response)
            by_multiplier[multiplier].append(1 if decision == "OPEN" else 0)
            margins_by_multiplier[multiplier].append(margin)
            margins_by_case[case][multiplier] = margin
            results.append(
                {
                    "case": case,
                    "multiplier": multiplier,
                    "seed": seed,
                    "decision": decision,
                    "decision_margin": margin,
                    "degenerated": is_degenerated(response),
                    "response": response,
                }
            )
    summary = {
        f"{multiplier:+.2f}": {
            "open_rate": sum(values) / len(values),
            "mean_decision_margin": sum(margins_by_multiplier[multiplier])
            / len(margins_by_multiplier[multiplier]),
            "samples": len(values),
        }
        for multiplier, values in by_multiplier.items()
    }
    directional_cases = sum(
        case_margins[-0.30] < case_margins[0.0] < case_margins[0.30]
        for case_margins in margins_by_case.values()
    )
    flipped_cases = sum(
        case_margins[-0.30] < 0 <= case_margins[0.30]
        for case_margins in margins_by_case.values()
    )
    payload = {
        "model": backend.name,
        "multipliers": MULTIPLIERS,
        "summary": summary,
        "direction_test": {
            "ordered_negative_neutral_positive": directional_cases,
            "decision_flips_negative_to_positive": flipped_cases,
            "cases": len(CASES),
        },
        "results": results,
    }
    output = ROOT / "artifacts" / "steering_evaluation.json"
    output.parent.mkdir(exist_ok=True)
    output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    return output, payload


if __name__ == "__main__":
    path, report = run_evaluation()
    print(path)
    print(json.dumps(report["summary"], indent=2))