640 lines
25 KiB
Python
640 lines
25 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Slice 19: tiny sequence-sensitive pragmatics specialists (binary action gate)
|
|
=============================================================================
|
|
|
|
A narrow binary specialist: executable request vs semantically similar
|
|
non-executable utterance, generalizing across surface-generator families.
|
|
|
|
Commands
|
|
grouped — grouped semantic CV (existing cv_fold), all arch/sizes, saves
|
|
per-config per-fold model checkpoints + OOF proba per variant
|
|
metrics — aggregate saved grouped-CV results into the report tables
|
|
(binary metrics, threshold curves, pair ordering, stress)
|
|
lfo — leave-generator-out. All sizes on capability_question; the
|
|
other present families for the leading config only.
|
|
e5baseline— frozen-e5 logistic + MLP baselines: grouped, cap-Q LOFO, pair
|
|
ordering (stress flagged NA — no re-embed on this box)
|
|
runtime — params / sizes / latency / tokenization for each candidate
|
|
|
|
Primary metrics (brief §2): cap-Q LOFO FA, pair ordering acc, pair margin,
|
|
grouped P/R, voice stress, fold variance. Aggregate accuracy is secondary.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import resource
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import warnings
|
|
from collections import Counter
|
|
|
|
import numpy as np
|
|
|
|
warnings.filterwarnings("ignore")
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from slice19_bpe import BpeVocab, CharVocab, normalize_match_text, strip_punct
|
|
|
|
EMBEDDING_PATH = "/tmp/mvn-experiment/embeddings.json"
|
|
RESULTS_DIR = "/tmp/mvn-s19"
|
|
SEED = 42
|
|
MAX_CHAR = 64
|
|
MAX_BPE = 24
|
|
BPE_VOCAB = 2048
|
|
|
|
PRESENT_FAMILIES = [
|
|
"polite_request",
|
|
"modal_request",
|
|
"first_person_request",
|
|
"reordered_target",
|
|
"capability_question",
|
|
"question",
|
|
]
|
|
ABSENT_FAMILIES = ["negation", "reported_speech", "quotation", "hypothetical"]
|
|
|
|
FAMILY_PRIORITY = [
|
|
"capability_question", "question", "first_person_request",
|
|
"modal_request", "polite_request", "reordered_target", "direct_imperative",
|
|
]
|
|
|
|
DEVICES = [
|
|
"свет", "люстру", "люстра", "жалюзи", "вытяжку", "вытяжка",
|
|
"вентилятор", "кондиционер", "телевизор", "лампу", "лампа",
|
|
"музыку", "музыка", "плеер", "колонку", "колонки", "чайник",
|
|
"бойлер", "обогреватель", "пылесос", "пылесосом", "пол",
|
|
"поливалки", "полив", "арка", "шторы", "штору",
|
|
"динамики", "дверь", "двери", "замок", "гараж", "ворота",
|
|
"кофе", "пасту", "зубы", "крючки", "лаймо", "куртку",
|
|
"будильник", "таймер", "напоминание", "расписание",
|
|
]
|
|
|
|
ARCH_CONFIGS = {
|
|
"char_cnn": ["tiny", "medium"],
|
|
"bigru": ["tiny", "medium", "large"],
|
|
"tiny_transformer": ["small", "medium"],
|
|
}
|
|
|
|
TRAIN_HYPER = {
|
|
"char_cnn": dict(epochs=20, lr=1e-3, bs=64, clip=None),
|
|
"bigru": dict(epochs=15, lr=5e-4, bs=64, clip=1.0),
|
|
"tiny_transformer": dict(epochs=25, lr=5e-4, bs=64, clip=1.0),
|
|
}
|
|
|
|
VARIANTS = ["orig", "nofinal", "strip"]
|
|
|
|
|
|
# ─── data ───────────────────────────────────────────────────────────────────
|
|
|
|
def load_dev():
|
|
with open(EMBEDDING_PATH) as f:
|
|
data = json.load(f)
|
|
rows = [e for e in data["examples"] if e["dev_pool"]]
|
|
out = []
|
|
for e in rows:
|
|
n_text = normalize_match_text(e["text"])
|
|
out.append({
|
|
"text_orig": n_text,
|
|
"text_nofinal": re.sub(r"[?.!,;:]+$", "", n_text),
|
|
"text_strip": strip_punct(n_text),
|
|
"route": e["route"],
|
|
"y": 1 if e["route"] == "action" else 0,
|
|
"cv_fold": e["cv_fold"],
|
|
"split_group": e["split_group"],
|
|
"tags": set(e.get("tags", [])),
|
|
"source_id": e["source_id"],
|
|
"emb": np.asarray(e["embedding"], dtype=np.float32),
|
|
})
|
|
return out
|
|
|
|
|
|
def family_of(tags):
|
|
for f in FAMILY_PRIORITY:
|
|
if f in tags:
|
|
return f
|
|
return None
|
|
|
|
|
|
def build_pairs(rows, texts):
|
|
"""Capability-question vs action-by-shared-object pairs (eval only)."""
|
|
def object_nouns(t):
|
|
found = set()
|
|
tl = t.lower()
|
|
for d in DEVICES:
|
|
if re.search(r"\b" + re.escape(d), tl):
|
|
found.add(d)
|
|
return found
|
|
|
|
pairs = []
|
|
for cidx, r in enumerate(rows):
|
|
if "capability_question" not in r["tags"]:
|
|
continue
|
|
c_obj = object_nouns(texts[cidx])
|
|
if not c_obj:
|
|
continue
|
|
dom = re.search(r"kq-cap-([^-]+)", r["source_id"])
|
|
dom = dom.group(1) if dom else None
|
|
for aidx, ra in enumerate(rows):
|
|
if ra["y"] != 1:
|
|
continue
|
|
a_obj = object_nouns(texts[aidx])
|
|
if not (c_obj & a_obj):
|
|
continue
|
|
a_dom = re.search(r"^([a-z]+)-", ra["source_id"])
|
|
a_dom = a_dom.group(1) if a_dom else None
|
|
if dom and a_dom and dom != a_dom:
|
|
continue
|
|
pairs.append((cidx, aidx))
|
|
return pairs
|
|
|
|
|
|
def pair_metrics(pairs, proba):
|
|
if not pairs:
|
|
return {"pairs": 0}
|
|
margins = []
|
|
ties = 0
|
|
order = 0
|
|
rev = 0
|
|
for cidx, aidx in pairs:
|
|
pc, pa = proba[cidx], proba[aidx]
|
|
margins.append(pa - pc)
|
|
if pa > pc:
|
|
order += 1
|
|
elif pa == pc:
|
|
ties += 1
|
|
else:
|
|
rev += 1
|
|
m = np.array(margins)
|
|
return {
|
|
"pairs": len(pairs),
|
|
"ordering_acc": order / len(pairs),
|
|
"mean_margin": float(m.mean()),
|
|
"median_margin": float(np.median(m)),
|
|
"ties": ties,
|
|
"reversed": rev,
|
|
}
|
|
|
|
|
|
# ─── tokenizers ─────────────────────────────────────────────────────────────
|
|
|
|
def build_tokenizers(rows):
|
|
nat_texts = [r["text_orig"] for r in rows]
|
|
char_vocab = CharVocab(nat_texts)
|
|
bpe = BpeVocab(nat_texts, vocab_size=BPE_VOCAB)
|
|
return char_vocab, bpe
|
|
|
|
|
|
def encode_all(rows, tokenizer, kind):
|
|
"""Return dict variant -> (N, max_len) int64 array."""
|
|
max_len = MAX_CHAR if kind == "char" else MAX_BPE
|
|
out = {}
|
|
for v in VARIANTS:
|
|
arr = np.zeros((len(rows), max_len), dtype=np.int64)
|
|
for i, r in enumerate(rows):
|
|
if kind == "char":
|
|
ids = tokenizer.encode(r[f"text_{v}"], max_len)
|
|
arr[i, :len(ids)] = ids
|
|
else:
|
|
ids = tokenizer.encode(r[f"text_{v}"])[:max_len]
|
|
arr[i, :len(ids)] = ids
|
|
out[v] = arr
|
|
return out
|
|
|
|
|
|
# ─── training ───────────────────────────────────────────────────────────────
|
|
|
|
def build_model(arch, size, vocab_size, max_len):
|
|
import torch
|
|
from slice19_models import CharCNN, BiGRU, TinyTransformer
|
|
if arch == "char_cnn":
|
|
c = (dict(embed_dim=32, filters=64, widths=[3, 4, 5]) if size == "tiny"
|
|
else dict(embed_dim=64, filters=160, widths=[2, 3, 4, 5]))
|
|
return CharCNN(vocab_size, c["embed_dim"], c["filters"], c["widths"])
|
|
if arch == "bigru":
|
|
c = (dict(embed_dim=64, hidden=64) if size == "tiny" else
|
|
(dict(embed_dim=128, hidden=128) if size == "medium" else
|
|
dict(embed_dim=256, hidden=256)))
|
|
return BiGRU(vocab_size, c["embed_dim"], c["hidden"])
|
|
c = (dict(d_model=128, n_layers=2, n_heads=4) if size == "small" else
|
|
dict(d_model=192, n_layers=4, n_heads=4))
|
|
return TinyTransformer(vocab_size, c["d_model"], c["n_layers"], c["n_heads"],
|
|
max_len=max_len)
|
|
|
|
|
|
def train_binary(X, y, arch, size, vocab_size, seed_offset=0, log=False):
|
|
import torch
|
|
torch.manual_seed(SEED + seed_offset)
|
|
np.random.seed(SEED + seed_offset)
|
|
Xt = torch.from_numpy(X)
|
|
yt = torch.from_numpy(y.astype(np.float32))
|
|
model = build_model(arch, size, vocab_size, X.shape[1])
|
|
h = TRAIN_HYPER[arch]
|
|
opt = torch.optim.AdamW(model.parameters(), lr=h["lr"], weight_decay=1e-4)
|
|
lossf = torch.nn.BCEWithLogitsLoss()
|
|
n = X.shape[0]
|
|
model.train()
|
|
t0 = time.time()
|
|
for epoch in range(h["epochs"]):
|
|
perm = torch.randperm(n)
|
|
running = 0.0
|
|
n_b = 0
|
|
for start in range(0, n, h["bs"]):
|
|
idx = perm[start:start + h["bs"]]
|
|
xb = Xt[idx]
|
|
if xb.dim() == 1:
|
|
xb = xb.unsqueeze(0)
|
|
logits = model(xb)
|
|
loss = lossf(logits, yt[idx])
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
if h["clip"]:
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), h["clip"])
|
|
opt.step()
|
|
running += float(loss)
|
|
n_b += 1
|
|
if log and (epoch + 1) % 5 == 0:
|
|
print(f" epoch {epoch+1}/{h['epochs']} loss {running/max(n_b,1):.4f}")
|
|
return model, time.time() - t0
|
|
|
|
|
|
def predict_proba(model, X, bs=256):
|
|
import torch
|
|
model.eval()
|
|
out = []
|
|
with torch.no_grad():
|
|
Xt = torch.from_numpy(X)
|
|
for start in range(0, X.shape[0], bs):
|
|
xb = Xt[start:start + bs]
|
|
if xb.dim() == 1:
|
|
xb = xb.unsqueeze(0)
|
|
logits = model(xb)
|
|
out.append(torch.sigmoid(logits).numpy())
|
|
return np.concatenate(out)
|
|
|
|
|
|
# ─── metrics helpers ────────────────────────────────────────────────────────
|
|
|
|
def binary_metrics(yt, proba, thr=0.5):
|
|
yp = (proba >= thr).astype(int)
|
|
tp = int(((yt == 1) & (yp == 1)).sum())
|
|
fp = int(((yt == 0) & (yp == 1)).sum())
|
|
fn = int(((yt == 1) & (yp == 0)).sum())
|
|
from sklearn.metrics import roc_auc_score, average_precision_score
|
|
roc = roc_auc_score(yt, proba) if len(np.unique(yt)) > 1 else 0.0
|
|
pr = average_precision_score(yt, proba) if len(np.unique(yt)) > 1 else 0.0
|
|
return {
|
|
"n": int(len(yt)), "tp": tp, "fp": fp, "fn": fn,
|
|
"P": tp / max(tp + fp, 1), "R": tp / max(tp + fn, 1),
|
|
"FA": fp, "FA_rate": fp / max(len(yt), 1),
|
|
"ROC_AUC": float(roc), "PR_AUC": float(pr),
|
|
}
|
|
|
|
|
|
def threshold_curve(yt, proba, thr_grid):
|
|
rows = []
|
|
for thr in thr_grid:
|
|
yp = (proba >= thr).astype(int)
|
|
tp = int(((yt == 1) & (yp == 1)).sum())
|
|
fp = int(((yt == 0) & (yp == 1)).sum())
|
|
fn = int(((yt == 1) & (yp == 0)).sum())
|
|
P = tp / max(tp + fp, 1)
|
|
R = tp / max(tp + fn, 1)
|
|
rows.append({
|
|
"thr": round(float(thr), 4), "P": round(P, 4), "R": round(R, 4),
|
|
"FA": fp, "FA_rate": round(fp / max(len(yt), 1), 4),
|
|
})
|
|
return rows
|
|
|
|
|
|
def operating_points(rows, pair_proba, pairs):
|
|
"""Report P>=0.95/0.98/0.99 points with pair separation at threshold."""
|
|
res = {}
|
|
for target in (0.95, 0.98, 0.99):
|
|
pts = [r for r in rows if r["P"] >= target and r["R"] > 0.0]
|
|
if not pts:
|
|
res[str(target)] = None
|
|
continue
|
|
best = max(pts, key=lambda r: r["R"])
|
|
# pair separation at that operating point
|
|
sep = pair_sep_at(best["thr"], pair_proba, pairs)
|
|
best = dict(best); best["pair_sep"] = round(sep, 4)
|
|
res[str(target)] = best
|
|
return res
|
|
|
|
|
|
def pair_sep_at(thr, pair_proba, pairs):
|
|
"""Fraction of pairs where action>=thr and cap<thr."""
|
|
if not pairs:
|
|
return 0.0
|
|
ok = 0
|
|
for cidx, aidx in pairs:
|
|
if pair_proba[aidx] >= thr and pair_proba[cidx] < thr:
|
|
ok += 1
|
|
return ok / len(pairs)
|
|
|
|
|
|
def fold_variance(fold_rows):
|
|
return {
|
|
"folds": [
|
|
{
|
|
"fold": fr["fold"],
|
|
"ROC_AUC": fr["ROC_AUC"], "PR_AUC": fr["PR_AUC"],
|
|
"P": fr["P"], "R": fr["R"], "FA": fr["FA"], "n": fr["n"],
|
|
}
|
|
for fr in fold_rows
|
|
]
|
|
}
|
|
|
|
|
|
# ─── subcommands ────────────────────────────────────────────────────────────
|
|
|
|
def _result_path():
|
|
os.makedirs(RESULTS_DIR, exist_ok=True)
|
|
return RESULTS_DIR
|
|
|
|
|
|
def cmd_grouped(args):
|
|
rows = load_dev()
|
|
char_vocab, bpe = build_tokenizers(rows)
|
|
y = np.array([r["y"] for r in rows])
|
|
folds = np.array([r["cv_fold"] for r in rows])
|
|
X_char = encode_all(rows, char_vocab, "char")
|
|
X_bpe = encode_all(rows, bpe, "bpe")
|
|
tokenizers = {"char": char_vocab, "bpe": bpe}
|
|
|
|
os.makedirs(RESULTS_DIR, exist_ok=True)
|
|
# save tokenizer metadata for reproducibility
|
|
meta = {
|
|
"char_vocab": char_vocab.size,
|
|
"char_vocab_sample": char_vocab.id_to_char[:50],
|
|
"bpe_vocab": bpe.size,
|
|
"bpe_serialized_bytes": bpe.serialized_bytes(),
|
|
"max_char": MAX_CHAR, "max_bpe": MAX_BPE,
|
|
"n": len(rows),
|
|
}
|
|
with open(os.path.join(RESULTS_DIR, "corpus_meta.json"), "w") as f:
|
|
json.dump(meta, f)
|
|
|
|
kind_of = {"char_cnn": "char", "bigru": "bpe", "tiny_transformer": "bpe"}
|
|
for arch, sizes in ARCH_CONFIGS.items():
|
|
kind = kind_of[arch]
|
|
X = X_char if kind == "char" else X_bpe
|
|
for size in sizes:
|
|
name = f"{arch}_{size}"
|
|
os.makedirs(os.path.join(RESULTS_DIR, "models"), exist_ok=True)
|
|
probs = {v: np.zeros(len(rows)) for v in VARIANTS}
|
|
fold_rows = []
|
|
for fold in range(5):
|
|
tr = folds != fold
|
|
te = folds == fold
|
|
Xtr = X["strip"][tr]
|
|
model, secs = train_binary(Xtr, y[tr], arch, size,
|
|
vocab_size=tokenizers[kind].size,
|
|
seed_offset=fold, log=args.verbose)
|
|
torch_models_dir = os.path.join(RESULTS_DIR, "models")
|
|
import torch
|
|
torch.save(model.state_dict(), os.path.join(torch_models_dir, f"{name}_fold{fold}.pt"))
|
|
for v in VARIANTS:
|
|
probs[v][te] = predict_proba(model, X[v][te])
|
|
fold_m = binary_metrics(y[te], probs["strip"][te])
|
|
fold_m["fold"] = fold
|
|
fold_rows.append(fold_m)
|
|
print(f" {name} fold {fold}: ROC={fold_m['ROC_AUC']:.3f} "
|
|
f"PR={fold_m['PR_AUC']:.3f} P={fold_m['P']:.3f} R={fold_m['R']:.3f} "
|
|
f"FA={fold_m['FA']} n={fold_m['n']} ({secs:.1f}s)")
|
|
np.savez(os.path.join(RESULTS_DIR, f"{name}_probs.npz"),
|
|
var_orig=probs["orig"], var_nofinal=probs["nofinal"],
|
|
var_strip=probs["strip"])
|
|
summary = binary_metrics(y, probs["strip"])
|
|
print(f" {name} OOF: ROC={summary['ROC_AUC']:.3f} PR={summary['PR_AUC']:.3f} "
|
|
f"P={summary['P']:.3f} R={summary['R']:.3f} FA={summary['FA']}")
|
|
print("grouped done")
|
|
|
|
|
|
def cmd_metrics(args):
|
|
rows = load_dev()
|
|
y = np.array([r["y"] for r in rows])
|
|
tags = [r["tags"] for r in rows]
|
|
pairs = build_pairs(rows, [r["text_strip"] for r in rows])
|
|
print(f"pairs={len(pairs)}")
|
|
|
|
out = {}
|
|
for arch, sizes in ARCH_CONFIGS.items():
|
|
for size in sizes:
|
|
name = f"{arch}_{size}"
|
|
fp = os.path.join(RESULTS_DIR, f"{name}_probs.npz")
|
|
if not os.path.exists(fp):
|
|
continue
|
|
z = np.load(fp)
|
|
entry = {"name": name, "arch": arch, "size": size}
|
|
# OOF binary on strip variant (primary training input)
|
|
entry["strip"] = binary_metrics(y, z["var_strip"])
|
|
entry["pairs"] = {}
|
|
entry["pairs"]["strip"] = pair_metrics(pairs, z["var_strip"])
|
|
entry["pairs"]["orig"] = pair_metrics(pairs, z["var_orig"])
|
|
entry["pairs"]["nofinal"] = pair_metrics(pairs, z["var_nofinal"])
|
|
# stress: same OOF models, per-variant metrics
|
|
entry["stress"] = {}
|
|
for v in VARIANTS:
|
|
p = z[f"var_{v}"]
|
|
entry["stress"][v] = {
|
|
"all_FA": binary_metrics(y, p)["FA_rate"],
|
|
"capQ_FA": capq_fa(tags, y, p),
|
|
"modal_recall": modal_recall(tags, y, p),
|
|
}
|
|
# threshold curve + operating points on strip
|
|
entry["curve"] = threshold_curve(y, z["var_strip"],
|
|
np.arange(0.30, 1.0, 0.02))
|
|
entry["ops"] = operating_points(entry["curve"], z["var_strip"], pairs)
|
|
out[name] = entry
|
|
|
|
with open(os.path.join(RESULTS_DIR, "metrics.json"), "w") as f:
|
|
json.dump(out, f, indent=2, default=str)
|
|
print(json.dumps(out, indent=2, default=str))
|
|
|
|
|
|
def capq_fa(tags, y, proba):
|
|
mask = np.array(["capability_question" in t for t in tags])
|
|
if mask.sum() == 0:
|
|
return 0.0
|
|
sub = proba[mask]
|
|
return float((sub >= 0.5).sum() / mask.sum())
|
|
|
|
|
|
def modal_recall(tags, y, proba):
|
|
mask = np.array([
|
|
(y[i] == 1 and ("polite_request" in tags[i] or "modal_request" in tags[i]))
|
|
for i in range(len(y))
|
|
])
|
|
if mask.sum() == 0:
|
|
return 0.0
|
|
sub = proba[mask]
|
|
return float((sub >= 0.5).sum() / mask.sum())
|
|
|
|
|
|
def cmd_lfo(args):
|
|
rows = load_dev()
|
|
char_vocab, bpe = build_tokenizers(rows)
|
|
y = np.array([r["y"] for r in rows])
|
|
tags = [r["tags"] for r in rows]
|
|
X_char = encode_all(rows, char_vocab, "char")["strip"]
|
|
X_bpe = encode_all(rows, bpe, "bpe")["strip"]
|
|
kind_of = {"char_cnn": "char", "bigru": "bpe", "tiny_transformer": "bpe"}
|
|
|
|
results = {}
|
|
tokenizers = {"char": char_vocab, "bpe": bpe}
|
|
# capability_question LOFO for every config/size (the critical case)
|
|
for arch, sizes in ARCH_CONFIGS.items():
|
|
X = X_char if kind_of[arch] == "char" else X_bpe
|
|
src_idx = np.array(["capability_question" not in t for t in tags])
|
|
tgt_idx = np.array(["capability_question" in t for t in tags])
|
|
for size in sizes:
|
|
model, _ = train_binary(X[src_idx], y[src_idx], arch, size,
|
|
vocab_size=tokenizers[kind_of[arch]].size,
|
|
seed_offset=17)
|
|
p = predict_proba(model, X[tgt_idx])
|
|
m = binary_metrics(y[tgt_idx], p)
|
|
results[f"{arch}_{size}:capability_question"] = m
|
|
print(f"LFO ability {arch}_{size}: cap rows={m['n']} "
|
|
f"pos={int(y[tgt_idx].sum())} FA={m['FA']} FA_rate={m['FA_rate']:.3f} "
|
|
f"P={m['P']:.3f} R={m['R']:.3f} acc={1-m['FA_rate']:.3f}")
|
|
|
|
# full family LOFO for the leading config per architecture
|
|
leading = {"char_cnn": "char_cnn_medium", "bigru": "bigru_tiny",
|
|
"tiny_transformer": "tiny_transformer_small"}
|
|
for arch, name in leading.items():
|
|
X = X_char if kind_of[arch] == "char" else X_bpe
|
|
for fam in PRESENT_FAMILIES:
|
|
src = np.array([fam not in t for t in tags])
|
|
tgt = np.array([fam in t for t in tags])
|
|
model, _ = train_binary(X[src], y[src], arch, name.split("_")[-1],
|
|
vocab_size=tokenizers[kind_of[arch]].size,
|
|
seed_offset=41)
|
|
p = predict_proba(model, X[tgt])
|
|
m = binary_metrics(y[tgt], p)
|
|
results[f"{name}:{fam}"] = m
|
|
print(f"LFO {fam}: {name} rows={m['n']} pos={int(y[tgt].sum())} "
|
|
f"P={m['P']:.3f} R={m['R']:.3f} FA={m['FA']} acc={1-m['FA_rate']:.3f}")
|
|
|
|
with open(os.path.join(RESULTS_DIR, "lfo.json"), "w") as f:
|
|
json.dump(results, f, indent=2, default=str)
|
|
print("lfo done")
|
|
|
|
|
|
def cmd_e5baseline(args):
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.neural_network import MLPClassifier
|
|
rows = load_dev()
|
|
y = np.array([r["y"] for r in rows])
|
|
folds = np.array([r["cv_fold"] for r in rows])
|
|
X = np.vstack([r["emb"] for r in rows])
|
|
tags = [r["tags"] for r in rows]
|
|
pairs = build_pairs(rows, [r["text_strip"] for r in rows])
|
|
|
|
out = {}
|
|
for model_name, model, extra in [
|
|
("e5_linear", LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42), {}),
|
|
("e5_mlp_h32", MLPClassifier(hidden_layer_sizes=(32,), activation="relu",
|
|
solver="adam", alpha=0.01, random_state=42,
|
|
max_iter=800, early_stopping=True,
|
|
validation_fraction=0.15, n_iter_no_change=10), {}),
|
|
]:
|
|
proba = np.zeros(len(rows))
|
|
for fold in range(5):
|
|
tr = folds != fold
|
|
te = folds == fold
|
|
m2 = type(model)(**{k: v for k, v in model.get_params().items()})
|
|
m2.fit(X[tr], y[tr])
|
|
proba[te] = m2.predict_proba(X[te])[:, 1]
|
|
entry = {
|
|
"grouped": binary_metrics(y, proba),
|
|
"pairs": pair_metrics(pairs, proba),
|
|
}
|
|
# cap-Q leave-generator-out (train without the family)
|
|
src = np.array(["capability_question" not in t for t in tags])
|
|
tgt = np.array(["capability_question" in t for t in tags])
|
|
m3 = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42) \
|
|
if model_name == "e5_linear" else \
|
|
MLPClassifier(hidden_layer_sizes=(32,), alpha=0.01, random_state=42, max_iter=800)
|
|
m3.fit(X[src], y[src])
|
|
p3 = m3.predict_proba(X[tgt])[:, 1]
|
|
entry["capq_lofo"] = binary_metrics(y[tgt], p3)
|
|
entry["stress"] = "NA (no re-embed on this box)"
|
|
out[model_name] = entry
|
|
print(f"{model_name}: grouped FA_rate={entry['grouped']['FA_rate']:.4f} "
|
|
f"PR={entry['grouped']['PR_AUC']:.3f} capQ_LOFO_FA_rate={entry['capq_lofo']['FA_rate']:.4f} "
|
|
f"pairs={entry['pairs']['ordering_acc']:.3f}")
|
|
with open(os.path.join(RESULTS_DIR, "e5baseline.json"), "w") as f:
|
|
json.dump(out, f, indent=2, default=str)
|
|
print("e5baseline done")
|
|
|
|
|
|
def cmd_runtime(args):
|
|
import torch
|
|
rows = load_dev()
|
|
char_vocab, bpe = build_tokenizers(rows)
|
|
X_char = encode_all(rows, char_vocab, "char")["strip"]
|
|
X_bpe = encode_all(rows, bpe, "bpe")["strip"]
|
|
kind_of = {"char_cnn": "char", "bigru": "bpe", "tiny_transformer": "bpe"}
|
|
report = {}
|
|
for arch, sizes in ARCH_CONFIGS.items():
|
|
X = X_char if kind_of[arch] == "char" else X_bpe
|
|
for size in sizes:
|
|
name = f"{arch}_{size}"
|
|
model = build_model(arch, size, int(X.max()) + 1, X.shape[1])
|
|
n_params = sum(p.numel() for p in model.parameters())
|
|
fp32 = n_params * 4
|
|
int8 = n_params
|
|
model.eval()
|
|
# warmup + latency (batch-1, eval mode)
|
|
xb = torch.from_numpy(X[:1])
|
|
with torch.no_grad():
|
|
for _ in range(20):
|
|
model(xb)
|
|
# tokenization latency
|
|
if kind_of[arch] == "char":
|
|
t0 = time.perf_counter()
|
|
for r in rows[:1000]:
|
|
char_vocab.encode(r["text_strip"], MAX_CHAR)
|
|
tl = (time.perf_counter() - t0) / 1000
|
|
else:
|
|
t0 = time.perf_counter()
|
|
for r in rows[:1000]:
|
|
bpe.encode(r["text_strip"])
|
|
tl = (time.perf_counter() - t0) / 1000
|
|
lat = []
|
|
for _ in range(300):
|
|
t0 = time.perf_counter()
|
|
model(xb)
|
|
lat.append(time.perf_counter() - t0)
|
|
lat = np.array(lat) * 1e6
|
|
report[name] = {
|
|
"params": n_params, "fp32_bytes": fp32, "int8_bytes": int8,
|
|
"latency_us_mean": float(lat.mean()), "latency_us_p50": float(np.median(lat)),
|
|
"latency_us_p95": float(np.percentile(lat, 95)),
|
|
"throughput_b1": round(1e6 / float(lat.mean()), 1),
|
|
"tok_us": round(tl * 1e6, 1),
|
|
"tokenizer": "char" if kind_of[arch] == "char" else "bpe",
|
|
}
|
|
print(f"{name}: {n_params} params fp32={fp32/1024:.0f}KiB "
|
|
f"lat={lat.mean():.0f}us tok={tl*1e6:.1f}us")
|
|
with open(os.path.join(RESULTS_DIR, "runtime.json"), "w") as f:
|
|
json.dump(report, f, indent=2)
|
|
print("runtime done")
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("cmd", choices=["grouped", "metrics", "lfo", "e5baseline", "runtime"])
|
|
ap.add_argument("--verbose", action="store_true")
|
|
args = ap.parse_args()
|
|
t0 = time.time()
|
|
globals()[f"cmd_{args.cmd}"](args)
|
|
print(f"elapsed {time.time()-t0:.1f}s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |