whetstone_DSL/specialists/eval/eval_specialist.py

#!/usr/bin/env python3
"""
eval_specialist.py

Extended evaluation for whetstone specialist checkpoints.
Produces confusion matrix, per-class accuracy, confidence distribution,
and threshold analysis (accepted precision + abstain rate).

Output: JSON to stdout or --out file.

Usage:
    python3 specialists/eval/eval_specialist.py \
        --checkpoint /mnt/storage/fabricate_runs/whetstone_worker_type/checkpoint.bin \
        --db         /mnt/storage/fabricate_runs/whetstone_worker_type/fabricate.db \
        --dataset    specialists/data/generated/worker_type_eval.tsv \
        --labels     "implementer,reviewer,architect,qa" \
        [--out       specialists/eval/results/worker_type_eval.json]
        [--thresholds 0.5,0.7,0.8,0.9]
        [--seq-len   96]
"""

import argparse
import json
import sys
import time
from collections import defaultdict
from pathlib import Path

import numpy as np

FABRICATE_TOOLS = Path("/home/bill/Documents/WhetstoneAI_Fabricate/tools")
sys.path.insert(0, str(FABRICATE_TOOLS))

from eval_with_vocab import (
    load_checkpoint, load_vocab, load_layers, build_runtime,
    tokenize_batch, forward_batch,
)


def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def load_labels(s):
    p = Path(s)
    if p.exists():
        return [t.strip() for t in p.read_text().strip().split(",") if t.strip()]
    return [t.strip() for t in s.split(",") if t.strip()]


def evaluate(checkpoint_path, db_path, dataset_path, label_tokens,
             seq_len=96, batch_size=256):
    weights, _ = load_checkpoint(checkpoint_path)
    _, str_to_id = load_vocab(db_path)
    layers = load_layers(db_path)

    vocab_size = len(str_to_id)
    _, emb_cnt = layers["token_embedding"]
    D = emb_cnt // vocab_size
    _, ff1_cnt = layers["block0.ff1_weight"]
    F = ff1_cnt // D
    L = sum(1 for k in layers if k.startswith("block") and k.endswith(".qkv_weight"))
    cfg = dict(seq_len=seq_len, D=D, F=F, V=vocab_size, L=L)
    runtime = build_runtime(weights, layers, cfg)

    label_vocab_ids = []
    missing = []
    for tok in label_tokens:
        vid = str_to_id.get(tok)
        if vid is None:
            missing.append(tok)
            vid = 1
        label_vocab_ids.append(vid)
    label_vocab_ids_arr = np.array(label_vocab_ids, dtype=np.int32)

    labels_gt, texts = [], []
    with open(dataset_path) as f:
        for line in f:
            parts = line.strip().split("\t", 2)
            if len(parts) < 3:
                continue
            labels_gt.append(int(parts[0]))
            texts.append(parts[2])

    labels_gt = np.array(labels_gt, dtype=np.int32)
    n = len(labels_gt)
    n_labels = len(label_tokens)

    all_preds = []
    all_confs = []  # max probability for predicted class

    t0 = time.perf_counter()
    for start in range(0, n, batch_size):
        stop = min(start + batch_size, n)
        token_ids = tokenize_batch(texts[start:stop], str_to_id, seq_len)
        logits = forward_batch(token_ids, weights, layers, cfg, runtime=runtime)
        label_logits = logits[:, label_vocab_ids_arr]
        probs = softmax(label_logits.astype(np.float64))
        pred_idx = np.argmax(probs, axis=1)
        conf = probs[np.arange(len(pred_idx)), pred_idx]
        all_preds.append(pred_idx)
        all_confs.append(conf)

    eval_latency_ms = (time.perf_counter() - t0) * 1000 / max(n, 1)

    preds = np.concatenate(all_preds)
    confs = np.concatenate(all_confs)
    correct = (preds == labels_gt)

    # --- Overall accuracy ---
    overall_acc = float(correct.mean() * 100)

    # --- Per-class accuracy ---
    per_class = {}
    for i, name in enumerate(label_tokens):
        mask = labels_gt == i
        if mask.sum() == 0:
            per_class[name] = {"acc": None, "n": 0}
        else:
            per_class[name] = {
                "acc": float(correct[mask].mean() * 100),
                "n": int(mask.sum()),
            }

    # --- Confusion matrix ---
    confusion = np.zeros((n_labels, n_labels), dtype=int)
    for gt, pred in zip(labels_gt, preds):
        confusion[gt, pred] += 1

    # Adjacent confusion (common mistake pairs)
    confusion_summary = []
    for i in range(n_labels):
        for j in range(n_labels):
            if i != j and confusion[i, j] > 0:
                confusion_summary.append({
                    "gt": label_tokens[i],
                    "pred": label_tokens[j],
                    "count": int(confusion[i, j]),
                    "gt_total": int((labels_gt == i).sum()),
                })
    confusion_summary.sort(key=lambda x: -x["count"])

    # --- Confidence distribution ---
    conf_bins = [0.0, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 1.0]
    conf_hist = {}
    for lo, hi in zip(conf_bins[:-1], conf_bins[1:]):
        mask = (confs >= lo) & (confs < hi)
        n_in_bin = int(mask.sum())
        if n_in_bin == 0:
            continue
        acc_in_bin = float(correct[mask].mean() * 100) if n_in_bin > 0 else None
        conf_hist[f"{lo:.2f}-{hi:.2f}"] = {"n": n_in_bin, "acc": acc_in_bin}

    # --- Threshold analysis ---
    threshold_results = []
    for threshold in [0.5, 0.6, 0.7, 0.8, 0.9]:
        accepted_mask = confs >= threshold
        n_accepted = int(accepted_mask.sum())
        n_abstained = n - n_accepted
        if n_accepted == 0:
            threshold_results.append({
                "threshold": threshold,
                "accepted": 0,
                "abstain_rate": 1.0,
                "accepted_accuracy": None,
            })
        else:
            acc = float(correct[accepted_mask].mean() * 100)
            threshold_results.append({
                "threshold": threshold,
                "accepted": n_accepted,
                "abstain_rate": round(n_abstained / n, 3),
                "accepted_accuracy": round(acc, 1),
            })

    # --- Guardrail assessment ---
    # Find lowest threshold where accepted_accuracy >= 95%
    guardrail_threshold = None
    guardrail_abstain_rate = None
    for tr in threshold_results:
        if tr["accepted_accuracy"] is not None and tr["accepted_accuracy"] >= 95.0:
            guardrail_threshold = tr["threshold"]
            guardrail_abstain_rate = tr["abstain_rate"]
            break

    # --- Deployment verdict ---
    if overall_acc >= 97:
        verdict = "deploy_direct"
    elif guardrail_threshold is not None and guardrail_abstain_rate is not None and guardrail_abstain_rate <= 0.30:
        verdict = "deploy_with_guardrails"
    elif overall_acc >= 80:
        verdict = "scale_more"
    elif len(confusion_summary) > 0 and confusion_summary[0]["count"] / n > 0.10:
        verdict = "decompose_or_relabel"
    else:
        verdict = "needs_investigation"

    return {
        "overall_accuracy": round(overall_acc, 1),
        "n_eval": n,
        "per_class": per_class,
        "confusion_summary": confusion_summary[:10],
        "confidence_histogram": conf_hist,
        "threshold_analysis": threshold_results,
        "guardrail_threshold": guardrail_threshold,
        "guardrail_abstain_rate": guardrail_abstain_rate,
        "eval_latency_ms_per_example": round(eval_latency_ms, 3),
        "missing_label_tokens": missing,
        "model_config": {"D": D, "F": F, "L": L, "vocab": vocab_size},
        "verdict": verdict,
    }


def print_report(result, labels):
    print(f"\n{'='*60}")
    print(f"  Overall accuracy: {result['overall_accuracy']:.1f}%  (n={result['n_eval']})")
    print(f"  Verdict: {result['verdict']}")
    print(f"{'='*60}")

    print("\nPer-class accuracy:")
    for name, stat in result["per_class"].items():
        if stat["acc"] is None:
            print(f"  {name:20s}  (no examples)")
        else:
            print(f"  {name:20s}  {stat['acc']:5.1f}%  (n={stat['n']})")

    if result["confusion_summary"]:
        print("\nTop confusions (gt → pred):")
        for c in result["confusion_summary"][:6]:
            pct = c["count"] / c["gt_total"] * 100
            print(f"  {c['gt']:20s} → {c['pred']:20s}  {c['count']}x  ({pct:.0f}% of gt class)")

    print("\nThreshold analysis:")
    print(f"  {'thresh':>6}  {'accepted':>8}  {'abstain%':>8}  {'acc@thresh':>10}")
    for tr in result["threshold_analysis"]:
        abstain_pct = tr["abstain_rate"] * 100
        acc = f"{tr['accepted_accuracy']:.1f}%" if tr["accepted_accuracy"] is not None else "  n/a"
        print(f"  {tr['threshold']:>6.2f}  {tr['accepted']:>8}  {abstain_pct:>7.1f}%  {acc:>10}")

    if result["guardrail_threshold"]:
        print(f"\n  ✓ Guarded deployment viable at threshold={result['guardrail_threshold']:.2f}"
              f"  (abstain {result['guardrail_abstain_rate']*100:.0f}%)")
    else:
        print("\n  ✗ No threshold achieves 95% accepted accuracy")

    print(f"\nEval latency: {result['eval_latency_ms_per_example']:.3f} ms/example")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--db",         required=True)
    ap.add_argument("--dataset",    required=True)
    ap.add_argument("--labels",     required=True)
    ap.add_argument("--seq-len",    type=int, default=96)
    ap.add_argument("--batch-size", type=int, default=256)
    ap.add_argument("--out",        default=None, help="Write JSON result to file")
    ap.add_argument("--quiet",      action="store_true", help="Suppress human-readable output")
    args = ap.parse_args()

    label_tokens = load_labels(args.labels)

    result = evaluate(
        args.checkpoint, args.db, args.dataset,
        label_tokens, seq_len=args.seq_len, batch_size=args.batch_size,
    )

    if not args.quiet:
        print_report(result, label_tokens)

    if args.out:
        out_path = Path(args.out)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json.dumps(result, indent=2))
        print(f"\nJSON → {args.out}")

    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()