specialists: switch training to PyTorch, add capacity sweep infrastructure
- 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>
This commit is contained in:
159
specialists/scripts/infer.py
Normal file
159
specialists/scripts/infer.py
Normal file
@@ -0,0 +1,159 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
infer.py — Run a trained Whetstone specialist on text input.
|
||||
|
||||
Supports both binary specialists (true/false) and N-way classifiers.
|
||||
|
||||
Usage:
|
||||
# Single query
|
||||
python3 specialists/scripts/infer.py \
|
||||
--run-dir /mnt/storage/fabricate_runs/whetstone_verification_type \
|
||||
--labels "unit,integration,schema,smoke,docs" \
|
||||
--text "Sprint 200 integration summary + regression"
|
||||
|
||||
# Batch from file (one text per line)
|
||||
python3 specialists/scripts/infer.py \
|
||||
--run-dir /mnt/storage/fabricate_runs/whetstone_verification_type \
|
||||
--labels "unit,integration,schema,smoke,docs" \
|
||||
--input-file my_tasks.txt
|
||||
|
||||
# Interactive REPL
|
||||
python3 specialists/scripts/infer.py \
|
||||
--run-dir /mnt/storage/fabricate_runs/whetstone_verification_type \
|
||||
--labels "unit,integration,schema,smoke,docs"
|
||||
"""
|
||||
import sys
|
||||
import argparse
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
FABRICATE_TOOLS = Path("/home/bill/Documents/WhetstoneAI_Fabricate/tools")
|
||||
|
||||
|
||||
def load_specialist(run_dir: str, labels_str: str | None):
|
||||
"""Load checkpoint, vocab, and labels. Return inference-ready state."""
|
||||
sys.path.insert(0, str(FABRICATE_TOOLS))
|
||||
from eval_with_vocab import load_checkpoint, load_vocab, load_layers, build_runtime
|
||||
|
||||
ck = Path(run_dir) / "checkpoint.bin"
|
||||
db = Path(run_dir) / "fabricate.db"
|
||||
if not ck.exists():
|
||||
raise FileNotFoundError(f"No checkpoint.bin in {run_dir}")
|
||||
if not db.exists():
|
||||
raise FileNotFoundError(f"No fabricate.db in {run_dir}")
|
||||
|
||||
weights, _ = load_checkpoint(str(ck))
|
||||
_, str_to_id = load_vocab(str(db))
|
||||
layers = load_layers(str(db))
|
||||
|
||||
V = len(str_to_id)
|
||||
_, emb_cnt = layers['token_embedding']
|
||||
D = emb_cnt // V
|
||||
_, 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=96, D=D, F=F, V=V, L=L)
|
||||
runtime = build_runtime(weights, layers, cfg)
|
||||
|
||||
label_list = None
|
||||
if labels_str:
|
||||
p = Path(labels_str)
|
||||
if p.exists():
|
||||
label_list = [l.strip() for l in p.read_text().split(',')]
|
||||
else:
|
||||
label_list = [l.strip() for l in labels_str.split(',')]
|
||||
|
||||
return weights, layers, cfg, runtime, str_to_id, label_list
|
||||
|
||||
|
||||
def predict(text: str, weights, layers, cfg, runtime, str_to_id, label_list) -> tuple[str, float]:
|
||||
"""Predict label for a single text. Returns (predicted_label, confidence)."""
|
||||
sys.path.insert(0, str(FABRICATE_TOOLS))
|
||||
from eval_with_vocab import tokenize_batch, forward_batch
|
||||
|
||||
token_ids = tokenize_batch([text], str_to_id, cfg['seq_len'])
|
||||
logits = forward_batch(token_ids, weights, layers, cfg, runtime=runtime) # (1, V)
|
||||
|
||||
if label_list:
|
||||
# N-way: restrict to label vocab IDs, take argmax
|
||||
label_ids = np.array([str_to_id.get(l, 0) for l in label_list], dtype=np.int32)
|
||||
label_logits = logits[0, label_ids]
|
||||
probs = np.exp(label_logits - label_logits.max())
|
||||
probs /= probs.sum()
|
||||
idx = int(np.argmax(probs))
|
||||
return label_list[idx], float(probs[idx])
|
||||
else:
|
||||
# Binary true/false
|
||||
true_id = str_to_id.get('true', 3)
|
||||
false_id = str_to_id.get('false', 4)
|
||||
true_logit = logits[0, true_id]
|
||||
false_logit = logits[0, false_id]
|
||||
prob_true = float(np.exp(true_logit) / (np.exp(true_logit) + np.exp(false_logit)))
|
||||
if prob_true >= 0.5:
|
||||
return "true", prob_true
|
||||
else:
|
||||
return "false", 1.0 - prob_true
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--run-dir", required=True,
|
||||
help="Directory with checkpoint.bin and fabricate.db")
|
||||
ap.add_argument("--labels", default=None,
|
||||
help="Comma-sep label list, e.g. 'unit,integration,schema,smoke,docs'")
|
||||
ap.add_argument("--text", default=None, help="Single text to classify")
|
||||
ap.add_argument("--input-file", default=None, help="File with one text per line")
|
||||
ap.add_argument("--eval-tsv", default=None,
|
||||
help="Evaluate accuracy on a TSV file (label<TAB>hops<TAB>text)")
|
||||
args = ap.parse_args()
|
||||
|
||||
print(f"Loading specialist from {args.run_dir}...", flush=True)
|
||||
state = load_specialist(args.run_dir, args.labels)
|
||||
weights, layers, cfg, runtime, str_to_id, label_list = state
|
||||
print(f"Loaded. V={cfg['V']} D={cfg['D']} L={cfg['L']} labels={label_list}")
|
||||
|
||||
if args.eval_tsv:
|
||||
correct = total = 0
|
||||
with open(args.eval_tsv) as f:
|
||||
for line in f:
|
||||
parts = line.strip().split('\t', 2)
|
||||
if len(parts) < 3:
|
||||
continue
|
||||
true_idx = int(parts[0])
|
||||
text = parts[2]
|
||||
true_label = label_list[true_idx] if label_list else str(true_idx)
|
||||
pred, conf = predict(text, *state[:-1], label_list)
|
||||
if pred == true_label:
|
||||
correct += 1
|
||||
total += 1
|
||||
print(f"Accuracy: {correct}/{total} = {100*correct/max(1,total):.1f}%")
|
||||
return
|
||||
|
||||
if args.input_file:
|
||||
with open(args.input_file) as f:
|
||||
lines = [l.strip() for l in f if l.strip()]
|
||||
for text in lines:
|
||||
pred, conf = predict(text, weights, layers, cfg, runtime, str_to_id, label_list)
|
||||
print(f"{pred:15s} {conf:.3f} {text}")
|
||||
return
|
||||
|
||||
if args.text:
|
||||
pred, conf = predict(args.text, weights, layers, cfg, runtime, str_to_id, label_list)
|
||||
print(f"{pred} ({conf:.3f})")
|
||||
return
|
||||
|
||||
# Interactive REPL
|
||||
print("Enter task descriptions (Ctrl-D to quit):")
|
||||
while True:
|
||||
try:
|
||||
text = input("> ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
break
|
||||
if not text:
|
||||
continue
|
||||
pred, conf = predict(text, weights, layers, cfg, runtime, str_to_id, label_list)
|
||||
print(f" → {pred} (conf={conf:.3f})")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user