run_step5.py

"""
run_step5.py — Phase 4 Step 5: score BF16 / Q8 / Q5 / Q4 on the 156-item
test set and print a comparison table.

Pipeline per model (largest → smallest):
  1. Launch llama-server.exe (LLAMA_SERVER constant below)
  2. Poll /health until 200 OK (60 s timeout)
  3. Delegate translation to the EXISTING run_gguf_eval.py (via subprocess)
     → writes results/gguf-{label}/{direction}_hypotheses.json
  4. Terminate server; poll until port 8081 is free before next model

Scoring reuses eval.py's functions imported directly (compute_bleu,
compute_comet, load_hypotheses, print_report).  COMET model is loaded once
and reused across all four models × two directions.

Usage:
    python run_step5.py                   # full inference + scoring
    python run_step5.py --skip-inference  # score already-existing hypothesis files
    python run_step5.py --skip-comet      # BLEU only (fast, no GPU needed)
    python run_step5.py --label q4        # single model only (inference + score)
"""

import argparse
import math
import os
import socket
import subprocess
import sys
import time

import requests

# ---------------------------------------------------------------------------
# CONFIG — edit these two constants if paths differ on your machine
# ---------------------------------------------------------------------------

LLAMA_SERVER  = r"C:\Users\Kaito Ishiguro\Documents\llama-b9716-bin\llama-server.exe"
N_GPU_LAYERS  = 99   # layers offloaded to GPU; set 0 for CPU-only (slow but safe)

PORT          = 8081
SERVER_URL    = f"http://localhost:{PORT}"
CTX_SIZE      = 512
RESULTS_ROOT  = "results"

# Model ladder — run and display largest → smallest
MODELS: list[tuple[str, str]] = [
    ("bf16", "models/cat-translate-0.8b-travel-ugm-bf16.gguf"),
    ("q8",   "models/cat-translate-0.8b-travel-Q8_0.gguf"),
    ("q5",   "models/cat-translate-0.8b-travel-Q5_K_M.gguf"),
    ("q4",   "models/cat-translate-0.8b-travel-Q4_K_M.gguf"),
]
MODEL_MAP: dict[str, str] = dict(MODELS)


# ---------------------------------------------------------------------------
# Server lifecycle
# ---------------------------------------------------------------------------

def _kill_any_on_port() -> None:
    """Force-kill any process already holding our port (Windows netstat approach)."""
    try:
        out = subprocess.run(
            ["netstat", "-ano"], capture_output=True, text=True
        ).stdout
        for line in out.splitlines():
            if f":{PORT}" in line and "LISTENING" in line:
                pid = line.strip().split()[-1]
                subprocess.run(["taskkill", "/F", "/PID", pid], capture_output=True)
                time.sleep(1)
    except Exception:
        pass


def launch_server(model_path: str) -> subprocess.Popen:
    cmd = [
        LLAMA_SERVER,
        "-m",       model_path,
        "--port",   str(PORT),
        "-c",       str(CTX_SIZE),
        "-ngl",     str(N_GPU_LAYERS),
        "--log-disable",
    ]
    print(f"  [server] {' '.join(cmd)}")
    # stdout/stderr inherit from parent so errors stay visible
    return subprocess.Popen(cmd)


def wait_server_ready(timeout: int = 60) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            r = requests.get(f"{SERVER_URL}/health", timeout=3)
            if r.status_code == 200:
                return True
        except requests.exceptions.RequestException:
            pass
        time.sleep(2)
    return False


def kill_server(proc: subprocess.Popen) -> None:
    print("  [server] terminating …", flush=True)
    proc.terminate()
    try:
        proc.wait(timeout=10)
    except subprocess.TimeoutExpired:
        print("  [server] kill (did not exit within 10 s)")
        proc.kill()
        proc.wait()


def wait_port_free(timeout: int = 30) -> bool:
    """Block until port stops accepting connections (post-server-kill)."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            with socket.create_connection(("localhost", PORT), timeout=1):
                time.sleep(1)   # still occupied
        except OSError:
            return True         # refused → free
    return False


# ---------------------------------------------------------------------------
# Inference phase: one model at a time
# ---------------------------------------------------------------------------

def run_inference(label: str, model_path: str) -> None:
    results_dir = f"{RESULTS_ROOT}/gguf-{label}"
    print(f"\n{'=' * 64}")
    print(f"  INFERENCE  [{label}]  →  {results_dir}")
    print(f"{'=' * 64}")

    _kill_any_on_port()
    proc = launch_server(model_path)

    try:
        print("  [server] waiting for /health …", flush=True)
        if not wait_server_ready(timeout=60):
            proc.kill()
            proc.wait()
            raise RuntimeError(
                f"llama-server did not respond within 60 s for [{label}].\n"
                f"  server : {LLAMA_SERVER}\n"
                f"  model  : {model_path}\n"
                "Check that both paths exist and the binary is the correct build."
            )
        print("  [server] ready.\n", flush=True)

        # Delegate translation to the EXISTING run_gguf_eval.py — not reimplemented here
        subprocess.run(
            [
                sys.executable, "run_gguf_eval.py",
                "--results-dir", results_dir,
                "--server-url",  SERVER_URL,
                "--direction",   "both",
            ],
            check=True,
        )

    finally:
        kill_server(proc)
        print("  [server] waiting for port to free …", flush=True)
        freed = wait_port_free(timeout=30)
        if not freed:
            print("  [WARN] port 8081 did not free within 30 s — continuing anyway.")
        else:
            print("  [server] port free.\n", flush=True)


# ---------------------------------------------------------------------------
# Scoring phase — import from eval.py, load COMET once, score all labels
# ---------------------------------------------------------------------------

def score_all(labels: list[str], skip_comet: bool) -> dict:
    """
    Import scoring functions from the EXISTING eval.py; load COMET once.

    Returns nested dict:
      scores[label][direction] = {
          "bleu":  {"overall": float, "by_domain": {domain: float}},
          "comet": {"overall": float, "by_domain": {domain: float}},
      }
    """
    # Import from existing eval.py — not reimplementing any scoring logic
    from eval import (
        load_hypotheses,
        compute_bleu,
        compute_comet,
        print_report,
        COMET_MODEL,
    )

    comet_model = None
    if not skip_comet:
        from comet import download_model, load_from_checkpoint
        print(f"\nLoading COMET ({COMET_MODEL}) — downloads ~1.5 GB on first run …")
        comet_model = load_from_checkpoint(download_model(COMET_MODEL))

    scores: dict = {}

    for label in labels:
        results_dir = f"{RESULTS_ROOT}/gguf-{label}"
        scores[label] = {}

        print(f"\n{'=' * 64}")
        print(f"  SCORING  [{label}]  ←  {results_dir}")
        print(f"{'=' * 64}")

        for direction in ("en2ja", "ja2en"):
            hyp_path = os.path.join(results_dir, f"{direction}_hypotheses.json")
            if not os.path.exists(hyp_path):
                raise FileNotFoundError(
                    f"Hypothesis file not found: {hyp_path}\n"
                    "Run without --skip-inference to generate it."
                )
            items = load_hypotheses(direction, results_dir=results_dir)
            bleu  = compute_bleu(items, direction)
            comet = (
                compute_comet(items, comet_model)
                if comet_model is not None
                else {"overall": float("nan"), "by_domain": {}}
            )
            # Mirror the same formatted output eval.py would print
            print_report(direction, bleu, comet, items)
            scores[label][direction] = {"bleu": bleu, "comet": comet}

    return scores


# ---------------------------------------------------------------------------
# Comparison table
# ---------------------------------------------------------------------------

def _file_mb(path: str) -> float:
    try:
        return os.path.getsize(path) / (1024 * 1024)
    except OSError:
        return float("nan")


def _avg(a: float, b: float) -> float:
    return float("nan") if (math.isnan(a) or math.isnan(b)) else (a + b) / 2


def _domain_comet(scores: dict, label: str, domain: str) -> float:
    c_en = scores[label]["en2ja"]["comet"]["by_domain"].get(domain, float("nan"))
    c_ja = scores[label]["ja2en"]["comet"]["by_domain"].get(domain, float("nan"))
    return _avg(c_en, c_ja)


def print_comparison(scores: dict, labels: list[str], skip_comet: bool) -> None:
    # Compute bf16 overall COMET for delta column
    def overall_comet(label: str) -> float:
        return _avg(
            scores[label]["en2ja"]["comet"]["overall"],
            scores[label]["ja2en"]["comet"]["overall"],
        )

    bf16_comet = overall_comet("bf16") if "bf16" in scores else float("nan")

    def fmt_comet(v: float, w: int = 10) -> str:
        return f"{v:{w}.4f}" if not math.isnan(v) else f"{'—':>{w}}"

    def fmt_delta(label: str, w: int = 10) -> str:
        if math.isnan(bf16_comet):
            return f"{'—':>{w}}"
        if label == "bf16":
            return f"{'±0.0000':>{w}}"
        d = overall_comet(label) - bf16_comet
        return f"{d:+{w}.4f}"

    sep = "─" * 102

    print(f"\n{sep}")
    print("  PHASE 4 STEP 5 — QUANT COMPARISON  (156-item test set)")
    if skip_comet:
        print("  NOTE: COMET skipped (--skip-comet).  Run without flag for full results.")
    print(sep)
    print(
        f"  {'Label':<6}  {'Size(MB)':>9}  {'COMET(all)':>10}  {'ΔCOMET':>10}  "
        f"{'food':>10}  {'emerg':>10}  {'BLEU en→ja':>12}  {'BLEU ja→en':>12}"
    )
    print(f"  {'─'*6}  {'─'*9}  {'─'*10}  {'─'*10}  {'─'*10}  {'─'*10}  {'─'*12}  {'─'*12}")

    for label in labels:
        if label not in scores:
            continue
        path   = MODEL_MAP[label]
        mb     = _file_mb(path)
        oc     = overall_comet(label)
        fc     = _domain_comet(scores, label, "food")
        ec     = _domain_comet(scores, label, "emergency")
        b_en   = scores[label]["en2ja"]["bleu"]["overall"]
        b_ja   = scores[label]["ja2en"]["bleu"]["overall"]

        mb_s   = f"{mb:9.1f}" if not math.isnan(mb) else f"{'?':>9}"

        print(
            f"  {label:<6}  {mb_s}  "
            f"{fmt_comet(oc)}  {fmt_delta(label)}  "
            f"{fmt_comet(fc)}  {fmt_comet(ec)}  "
            f"{b_en:>12.2f}  {b_ja:>12.2f}"
        )

    print(sep)
    print(
        "  COMET model : wmt22-comet-da.  "
        "food / emerg COMET = mean(en→ja, ja→en) domain scores.\n"
        "  BLEU en→ja  : char tokenizer.  "
        "BLEU ja→en : 13a tokenizer.  "
        "Tokenizers differ — do not compare across directions.\n"
        "  Primary signal: overall COMET.  BLEU shown for secondary reference only."
    )
    print(sep)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Phase 4 Step 5 — evaluate all GGUF quants on the 156-item test set"
    )
    parser.add_argument(
        "--skip-inference",
        action="store_true",
        help="Skip server + translation; score hypothesis files already on disk",
    )
    parser.add_argument(
        "--skip-comet",
        action="store_true",
        help="Skip COMET scoring (fast, BLEU only)",
    )
    parser.add_argument(
        "--label",
        choices=[lbl for lbl, _ in MODELS],
        help="Run a single model only (default: all four)",
    )
    args = parser.parse_args()

    sys.stdout.reconfigure(encoding="utf-8")

    # Determine which models to run
    targets = [(lbl, p) for lbl, p in MODELS if (args.label is None or lbl == args.label)]

    # --- Pre-flight checks ---
    if not args.skip_inference:
        if not os.path.exists(LLAMA_SERVER):
            raise FileNotFoundError(
                f"LLAMA_SERVER not found: {LLAMA_SERVER}\n"
                "Edit the LLAMA_SERVER constant at the top of this file."
            )
        for label, path in targets:
            if not os.path.exists(path):
                raise FileNotFoundError(
                    f"GGUF not found for [{label}]: {path}"
                )

    # --- Inference phase ---
    if not args.skip_inference:
        for label, model_path in targets:
            run_inference(label, model_path)

    # --- Scoring phase ---
    labels_to_score = [lbl for lbl, _ in targets]
    scores = score_all(labels_to_score, skip_comet=args.skip_comet)

    # --- Comparison table ---
    ordered_labels = [lbl for lbl, _ in MODELS if lbl in scores]
    print_comparison(scores, ordered_labels, skip_comet=args.skip_comet)


if __name__ == "__main__":
    main()