3fc9deac5b
- train_specialist_pt.py: pure PyTorch trainer replacing Fabricate binary. No SQLite/WAL — saves checkpoint.pt + history.json only. Fabricate was causing D-state IO blocking on the HDD due to continuous WAL writes. - eval_specialist_pt.py: PyTorch eval with confusion matrix, threshold analysis, guardrail assessment, and deployment verdict. - capacity_sweep.py: updated to call PyTorch scripts; TIERS now include heads field (4/6/8) since PyTorch supports multi-head unlike Fabricate. - prereq_op_binary_split.py: evaluate_combined rewritten in PyTorch. - Performance fix in train_specialist_pt.py: load training data onto GPU once and sample via torch.randint instead of DataLoader (eliminates Python/CPU overhead on tiny datasets). Fix applied, not yet timed. Next session: time the fix, then run sweep_all_gates.sh --pilot-only. See HANDOFF-2026-03-31.md for full context. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
297 lines
10 KiB
Python
297 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
eval_specialist_pt.py
|
|
|
|
Extended evaluation for PyTorch specialist checkpoints.
|
|
Drop-in replacement for eval_specialist.py, but loads checkpoint.pt instead
|
|
of checkpoint.bin + fabricate.db.
|
|
|
|
Output: JSON to stdout (and optionally --out file).
|
|
|
|
Usage:
|
|
python3 specialists/eval/eval_specialist_pt.py \
|
|
--checkpoint /tmp/runs/worker_type/checkpoint.pt \
|
|
--dataset specialists/data/generated/worker_type_eval.tsv \
|
|
--labels "implementer,reviewer,architect,qa" \
|
|
[--out specialists/eval/results/worker_type_eval.json] \
|
|
[--quiet]
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
|
|
# ── Copied from train_specialist_pt.py ───────────────────────────────────────
|
|
|
|
def _tokenize(text):
|
|
return text.lower().split()
|
|
|
|
|
|
class Vocab:
|
|
PAD, UNK = 0, 1
|
|
|
|
def __init__(self, str_to_id=None):
|
|
if str_to_id is not None:
|
|
self.str_to_id = str_to_id
|
|
self.id_to_str = [None] * len(str_to_id)
|
|
for s, i in str_to_id.items():
|
|
self.id_to_str[i] = s
|
|
else:
|
|
self.str_to_id = {"<pad>": 0, "<unk>": 1}
|
|
self.id_to_str = ["<pad>", "<unk>"]
|
|
|
|
def encode(self, text, seq_len):
|
|
ids = [self.str_to_id.get(t, self.UNK) for t in _tokenize(text)]
|
|
ids = ids[:seq_len]
|
|
ids += [self.PAD] * (seq_len - len(ids))
|
|
return ids
|
|
|
|
def __len__(self):
|
|
return len(self.id_to_str)
|
|
|
|
|
|
class SpecialistModel(nn.Module):
|
|
def __init__(self, vocab_size, hidden_dim, layers, heads, n_labels, seq_len):
|
|
super().__init__()
|
|
self.seq_len = seq_len
|
|
self.tok_emb = nn.Embedding(vocab_size, hidden_dim, padding_idx=0)
|
|
self.pos_emb = nn.Embedding(seq_len, hidden_dim)
|
|
encoder_layer = nn.TransformerEncoderLayer(
|
|
d_model=hidden_dim,
|
|
nhead=heads,
|
|
dim_feedforward=hidden_dim * 4,
|
|
dropout=0.0,
|
|
batch_first=True,
|
|
)
|
|
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=layers)
|
|
self.head = nn.Linear(hidden_dim, n_labels)
|
|
|
|
def forward(self, x):
|
|
pad_mask = (x == 0)
|
|
pos = torch.arange(x.size(1), device=x.device).unsqueeze(0)
|
|
h = self.tok_emb(x) + self.pos_emb(pos)
|
|
h = self.encoder(h, src_key_padding_mask=pad_mask)
|
|
lengths = (~pad_mask).float().sum(dim=1, keepdim=True).clamp(min=1)
|
|
h = (h * (~pad_mask).unsqueeze(-1).float()).sum(dim=1) / lengths
|
|
return self.head(h)
|
|
|
|
|
|
# ── Eval ──────────────────────────────────────────────────────────────────────
|
|
|
|
def softmax(x):
|
|
e = np.exp(x - x.max(axis=-1, keepdims=True))
|
|
return e / e.sum(axis=-1, keepdims=True)
|
|
|
|
|
|
def evaluate(checkpoint_path, dataset_path, label_tokens, seq_len=96, batch_size=256):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
ckpt = torch.load(checkpoint_path, map_location=device)
|
|
cfg = ckpt["config"]
|
|
|
|
# Reconstruct vocab from checkpoint
|
|
vocab_obj = ckpt.get("vocab_obj")
|
|
if vocab_obj is None:
|
|
vocab_obj = Vocab(str_to_id=ckpt["vocab"])
|
|
|
|
model = SpecialistModel(
|
|
vocab_size=cfg["vocab_size"],
|
|
hidden_dim=cfg["hidden_dim"],
|
|
layers=cfg["layers"],
|
|
heads=cfg["heads"],
|
|
n_labels=cfg["n_labels"],
|
|
seq_len=cfg["seq_len"],
|
|
).to(device)
|
|
model.load_state_dict(ckpt["model"])
|
|
model.eval()
|
|
|
|
# Load eval data
|
|
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])
|
|
|
|
n = len(labels_gt)
|
|
labels_gt_arr = np.array(labels_gt, dtype=np.int32)
|
|
|
|
# Run inference
|
|
all_preds = []
|
|
all_confs = []
|
|
t0 = time.perf_counter()
|
|
|
|
with torch.no_grad():
|
|
for start in range(0, n, batch_size):
|
|
batch_texts = texts[start:start + batch_size]
|
|
ids = [vocab_obj.encode(t, cfg["seq_len"]) for t in batch_texts]
|
|
x = torch.tensor(ids, dtype=torch.long, device=device)
|
|
logits = model(x).cpu().numpy().astype(np.float64)
|
|
probs = softmax(logits)
|
|
pred = np.argmax(probs, axis=1)
|
|
conf = probs[np.arange(len(pred)), pred]
|
|
all_preds.append(pred)
|
|
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_arr)
|
|
overall_acc = float(correct.mean() * 100)
|
|
|
|
# Per-class accuracy
|
|
per_class = {}
|
|
for i, name in enumerate(label_tokens):
|
|
mask = labels_gt_arr == 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
|
|
n_labels = len(label_tokens)
|
|
confusion = np.zeros((n_labels, n_labels), dtype=int)
|
|
for gt, pred in zip(labels_gt_arr, preds):
|
|
confusion[gt, pred] += 1
|
|
|
|
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_arr == i).sum()),
|
|
})
|
|
confusion_summary.sort(key=lambda x: -x["count"])
|
|
|
|
# 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 (lowest threshold with ≥95% accepted accuracy)
|
|
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 confusion_summary 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],
|
|
"threshold_analysis": threshold_results,
|
|
"guardrail_threshold": guardrail_threshold,
|
|
"guardrail_abstain_rate": guardrail_abstain_rate,
|
|
"eval_latency_ms_per_example": round(eval_latency_ms, 3),
|
|
"model_config": cfg,
|
|
"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("--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)
|
|
ap.add_argument("--quiet", action="store_true")
|
|
# Accept --db for compatibility with old callers (ignored)
|
|
ap.add_argument("--db", default=None)
|
|
args = ap.parse_args()
|
|
|
|
label_tokens = [l.strip() for l in args.labels.split(",")]
|
|
result = evaluate(args.checkpoint, 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()
|