benchmarks/gate_g2_g3.py

"""Gates G2/G3: Director Spanish fidelity and llama.cpp CPU speed."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import statistics
import sys
import time
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from lightloom.core.config import CONFIG  # noqa: E402
from lightloom.director.director import (  # noqa: E402
    generate_shot,
    generate_shot_transformers,
    load_director,
    load_director_transformers,
)

RESULTS_DIR = ROOT / "benchmarks" / "results"
G2_JSON = RESULTS_DIR / "g2.json"
G3_JSON = RESULTS_DIR / "g3.json"

GOLDENS = [
    ("EN", "Every night he climbed the spiral stairs, counting each step like a prayer."),
    ("EN", "Far away, a market woke under orange awnings and wet stone."),
    ("EN", "The child opened the brass box and found a moth made of blue glass."),
    ("EN", "Years later, the same road was buried under snow and silence."),
    ("EN", "At sunrise, the keeper saw a second lighthouse blinking back."),
    ("ES", "Cada noche subía la escalera en espiral, contando cada peldaño como una oración."),
    ("ES", "Lejos de allí, un mercado despertaba bajo toldos naranjas y piedra mojada."),
    ("ES", "La niña abrió la caja de latón y encontró una polilla de vidrio azul."),
    ("ES", "Años después, el mismo camino estaba enterrado bajo nieve y silencio."),
    ("ES", "Al amanecer, el farero vio otro faro parpadeando de vuelta."),
]


def _dry_run() -> tuple[dict[str, Any], dict[str, Any]]:
    rows = [
        {
            "lang": lang,
            "valid_json": True,
            "elapsed_s": 0.01,
            "tok_s": 1000.0,
            "heuristic_fidelity": "dry_run",
        }
        for lang, _ in GOLDENS
    ]
    g2 = {
        "schema_version": 1,
        "gate": "G2",
        "authoritative": False,
        "valid_json_rate": 1.0,
        "spanish_cases": 5,
        "decision": "dry_run",
        "cases": rows,
    }
    g3 = {
        "schema_version": 1,
        "gate": "G3",
        "authoritative": False,
        "p50_s_per_beat": 0.01,
        "p95_s_per_beat": 0.01,
        "p50_tok_s": 1000.0,
        "decision": "dry_run",
    }
    return g2, g3


def _percentile(values: list[float], q: float) -> float:
    ordered = sorted(values)
    idx = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * q)))
    return ordered[idx]


def _score_fidelity(lang: str, raw: str, image_prompt: str) -> str:
    if lang != "ES":
        return "not_applicable"
    text = image_prompt.casefold()
    checks = {
        "mercado": ("market", "awning", "orange"),
        "polilla": ("moth", "glass", "blue"),
        "nieve": ("snow", "road", "silence"),
        "farero": ("lighthouse", "keeper"),
        "escalera": ("stair", "spiral"),
    }
    raw_l = raw.casefold()
    for marker, expected in checks.items():
        if marker in raw_l:
            hits = sum(1 for word in expected if word in text)
            return "pass" if hits >= max(1, len(expected) - 1) else "fail"
    return "review"


def run(dry_run: bool = False) -> tuple[dict[str, Any], dict[str, Any]]:
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    if dry_run:
        g2, g3 = _dry_run()
    else:
        import os
        import platform

        backend = os.getenv("LIGHTLOOM_DIRECTOR_BACKEND") or CONFIG.director_backend
        if backend == "transformers":
            tokenizer, model = load_director_transformers()

            def director(beat: str, bid: int) -> Any:
                return generate_shot_transformers(tokenizer, model, beat, beat_id=bid)
        else:
            llm = load_director(n_threads=os.cpu_count())

            def director(beat: str, bid: int) -> Any:
                return generate_shot(llm, beat, beat_id=bid)

        rows: list[dict[str, Any]] = []
        for idx, (lang, beat) in enumerate(GOLDENS, start=1):
            start = time.perf_counter()
            try:
                shot, meta = director(beat, idx)
                elapsed_s = float(meta["elapsed_s"])
                row = {
                    "lang": lang,
                    "beat": beat,
                    "valid_json": True,
                    "elapsed_s": round(elapsed_s, 3),
                    "tok_s": round(meta["tok_s"], 3) if meta["tok_s"] else None,
                    "completion_tokens": meta["completion_tokens"],
                    "attempts": meta.get("attempts"),
                    "decision": shot.decision,
                    "image_prompt_en": shot.image_prompt_en,
                    "heuristic_fidelity": _score_fidelity(lang, beat, shot.image_prompt_en),
                }
            except Exception as exc:  # noqa: BLE001
                row = {
                    "lang": lang,
                    "beat": beat,
                    "valid_json": False,
                    "elapsed_s": round(time.perf_counter() - start, 3),
                    "error": f"{type(exc).__name__}: {exc}",
                    "heuristic_fidelity": "fail" if lang == "ES" else "not_applicable",
                }
            rows.append(row)
        times = [row["elapsed_s"] for row in rows if row.get("valid_json")]
        tok_s = [row["tok_s"] for row in rows if row.get("tok_s")]
        valid_json_rate = sum(1 for row in rows if row.get("valid_json")) / len(rows)
        es_rows = [row for row in rows if row["lang"] == "ES"]
        es_pass = sum(1 for row in es_rows if row.get("heuristic_fidelity") == "pass")
        g2_decision = "aya_toggle" if valid_json_rate == 1.0 and es_pass >= 4 else "aya_must"
        p50_beat = statistics.median(times) if times else None
        if backend == "transformers":
            g3_decision = "director_gpu_ok" if (p50_beat and p50_beat <= 2.5) else "director_gpu_slow"
        else:
            g3_decision = "director_cpu" if times and _percentile(times, 0.95) <= 2.5 else "director_gpu_fallback"
        hardware = {
            "profile": os.getenv("LIGHTLOOM_PROFILE", "local"),
            "platform": platform.platform(),
            "cpu_count": os.cpu_count(),
            "director_backend": backend,
        }
        g2 = {
            "schema_version": 1,
            "gate": "G2",
            "authoritative": os.getenv("LIGHTLOOM_PROFILE") == "space",
            "hardware_profile": hardware,
            "valid_json_rate": valid_json_rate,
            "spanish_fidelity_passes": es_pass,
            "spanish_cases": len(es_rows),
            "decision": g2_decision,
            "cases": rows,
        }
        g3 = {
            "schema_version": 1,
            "gate": "G3",
            "authoritative": os.getenv("LIGHTLOOM_PROFILE") == "space",
            "hardware_profile": hardware,
            "p50_s_per_beat": round(statistics.median(times), 3) if times else None,
            "p95_s_per_beat": round(_percentile(times, 0.95), 3) if times else None,
            "p50_tok_s": round(statistics.median(tok_s), 3) if tok_s else None,
            "decision": g3_decision,
            "cases": rows,
        }
    G2_JSON.write_text(json.dumps(g2, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    G3_JSON.write_text(json.dumps(g3, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return g2, g3


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    g2, g3 = run(dry_run=args.dry_run)
    print(f"G2 {g2['decision']} -> {G2_JSON}")
    print(f"G3 {g3['decision']} -> {G3_JSON}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())