oev 0.1.0__py3-none-any.whl

oev/__init__.py

@@ -0,0 +1,7 @@
+"""OEV: a small System One decision model.
+
+State + typed questions in, calibrated probability distributions out,
+one forward pass.
+"""
+
+__version__ = "0.1.0"

oev/benchmark.py

@@ -0,0 +1,70 @@
+import argparse
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from oev.dataset import OEVDataset, collate
+from oev.model import HFBackboneOEV
+from oev.evaluate import load_model, ece
+from oev.calibrate import fit_temperature_from_logits
+from oev.tokenizer_hf import HFTokenPacker
+
+
+def stack_ragged(rows):
+    kmax = max(r.numel() for r in rows)
+    out = torch.full((len(rows), kmax), -1e4)
+    for i, r in enumerate(rows):
+        out[i, : r.numel()] = r
+    return out
+
+
+def evaluate_benchmark(checkpoint, data_dir, calibrate=True, batch_size=64):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = load_model(checkpoint, device)
+    if isinstance(model, HFBackboneOEV):
+        packer = HFTokenPacker(model.cfg["backbone"])
+        max_len = model.cfg["max_len"]
+    else:
+        packer = None
+        max_len = model.cfg.max_len
+    dl = DataLoader(OEVDataset(f"{data_dir}/test.jsonl", max_len, packer=packer), batch_size=batch_size, collate_fn=collate)
+    confs, corr = [], []
+    val_logits, val_labels = [], []
+    if calibrate:
+        vdl = DataLoader(OEVDataset(f"{data_dir}/valid.jsonl", max_len, packer=packer), batch_size=batch_size, collate_fn=collate)
+        with torch.no_grad():
+            for batch in vdl:
+                logits = model(batch["ids"].to(device), batch["pad_mask"].to(device), batch["anchor_pos"].to(device)) + batch["logits_mask"].to(device)
+                keep = batch["anchor_valid"].sum(dim=1)
+                for i in range(logits.size(0)):
+                    row = logits[i][: keep[i]]
+                    if row.numel() >= 2:
+                        val_logits.append(row)
+                        val_labels.append(batch["labels"][i])
+        temperature = fit_temperature_from_logits(stack_ragged(val_logits), torch.stack(val_labels))
+    else:
+        temperature = 1.0
+    with torch.no_grad():
+        for batch in dl:
+            logits = model(batch["ids"].to(device), batch["pad_mask"].to(device), batch["anchor_pos"].to(device)) + batch["logits_mask"].to(device)
+            probs = F.softmax(logits / temperature, dim=-1)
+            conf, pred = probs.max(dim=-1)
+            for i in range(len(batch["labels"])):
+                confs.append(conf[i].item())
+                corr.append(int(pred[i].item() == batch["labels"][i].item()))
+    return {
+        "accuracy": sum(corr) / len(corr),
+        "n": len(corr),
+        "ece": ece(confs, corr),
+        "temperature": temperature,
+    }
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--data-dir", required=True)
+    p.add_argument("--no-calibrate", action="store_true")
+    args = p.parse_args()
+    r = evaluate_benchmark(args.checkpoint, args.data_dir, calibrate=not args.no_calibrate)
+    for k, v in r.items():
+        print(f"{k}: {v:.4f}" if isinstance(v, float) else f"{k}: {v}")

oev/benchmark_ext.py

@@ -0,0 +1,213 @@
+"""Extended benchmark metrics for typed-decisions: soft accuracy, Brier score,
+score MAE, and per-workflow / per-primitive breakdowns.
+
+Mirrors the metric set published in Laya's typed-decisions table so results
+are directly comparable column-for-column.
+
+Usage:
+    python -m oev.benchmark_ext --checkpoint CKPT --data-dir data/typed
+Ensembles:
+    python -m oev.benchmark_ext --ckpts a.pt,b.pt,c.pt --data-dir data/typed
+"""
+
+import argparse
+import json
+
+import torch
+import torch.nn.functional as F
+
+from oev.evaluate import load_model
+from oev.tokenizer_hf import HFTokenPacker
+
+
+def _predict_probs(model, packer, state, question, device, gamma=1.0):
+    pq = {
+        "name": question["name"],
+        "type": question["type"],
+        "instructions": question.get("instructions", question["type"]),
+        "options": question["options"],
+        "answer": question["answer"],
+    }
+    ids, anchors, label = packer.pack(state, pq, model.cfg["max_len"])
+    tids = torch.tensor([ids], device=device)
+    pmask = torch.zeros(1, len(ids), dtype=torch.bool, device=device)
+    apos = torch.tensor([anchors], device=device)
+    with torch.no_grad():
+        logits = model(tids, pmask, apos)
+    probs = torch.softmax(logits[0].float(), dim=-1)
+    if gamma != 1.0:
+        probs = probs.clamp_min(1e-9) ** gamma
+        probs = probs / probs.sum()
+    return probs, label, question["type"], question.get("target")
+
+
+def ece_metric(confs, corrs, n_bins=10):
+    """Expected Calibration Error: weighted |confidence - accuracy| over bins."""
+    if not confs:
+        return 0.0
+    bins = [[] for _ in range(n_bins)]
+    for c, o in zip(confs, corrs):
+        b = min(int(c * n_bins), n_bins - 1)
+        bins[b].append((c, o))
+    e = 0.0
+    for b in bins:
+        if not b:
+            continue
+        acc = sum(o for _, o in b) / len(b)
+        conf = sum(c for c, _ in b) / len(b)
+        e += (len(b) / len(confs)) * abs(conf - acc)
+    return e
+
+
+def evaluate_metrics(checkpoints, data_dir, device="cuda", gamma=1.0, latency=False):
+    if device == "cuda" and not torch.cuda.is_available():
+        device = "cpu"
+    models = [load_model(c, device) for c in checkpoints]
+    packers = [HFTokenPacker(m.cfg["backbone"]) for m in models]
+
+    rows = [json.loads(l) for l in open(f"{data_dir}/test.jsonl", encoding="utf-8")]
+
+    n = correct = 0
+    soft_acc_sum = 0.0
+    brier_sum = 0.0
+    score_mae_sum = 0.0
+    score_n = 0
+    confs = []
+    corrs = []
+    by_type = {}
+    by_domain = {}
+
+    with torch.no_grad():
+        for r in rows:
+            for q in r["questions"]:
+                probs_sum = None
+                label = None
+                target = None
+                qtype = None
+                for m, p in zip(models, packers):
+                    probs, label, qtype, target = _predict_probs(m, p, r["state"], q, device, gamma=gamma)
+                    probs_sum = probs if probs_sum is None else probs_sum + probs
+                probs = probs_sum / len(models)
+
+                n += 1
+                hit = int(probs.argmax().item() == label)
+                correct += hit
+                confs.append(probs.max().item())
+                corrs.append(hit)
+
+                # soft accuracy: the probability mass the model put on the gold answer
+                soft_acc_sum += probs[label].item()
+
+                # Brier score against the gold one-hot (lower is better)
+                oh = torch.zeros_like(probs)
+                oh[label] = 1.0
+                brier_sum += ((probs - oh) ** 2).sum().item()
+
+                # score MAE: |expected level - gold level| for score questions
+                if qtype == "score":
+                    levels = torch.arange(probs.numel(), dtype=torch.float32, device=probs.device)
+                    ev = (probs * levels).sum().item()
+                    score_mae_sum += abs(ev - label)
+                    score_n += 1
+
+                agg = by_type.setdefault(qtype, [0, 0, 0.0])
+                agg[0] += hit
+                agg[1] += 1
+                agg[2] += probs[label].item()
+
+                agg = by_domain.setdefault(r.get("domain", "typed"), [0, 0, 0.0])
+                agg[0] += hit
+                agg[1] += 1
+                agg[2] += probs[label].item()
+
+    result = {
+        "accuracy": correct / n,
+        "soft_acc": soft_acc_sum / n,
+        "brier": brier_sum / n,
+        "score_mae": (score_mae_sum / score_n) if score_n else None,
+        "ece": ece_metric(confs, corrs),
+        "n": n,
+    }
+    print(f"accuracy : {result['accuracy']:.4f}")
+    print(f"soft acc : {result['soft_acc']:.4f}")
+    print(f"brier    : {result['brier']:.4f}")
+    if result["score_mae"] is not None:
+        print(f"score MAE: {result['score_mae']:.4f}")
+    print(f"ece      : {result['ece']:.4f}")
+    print(f"n        : {result['n']}")
+
+    print("\nper primitive:")
+    for t, (c, tot, s) in sorted(by_type.items()):
+        print(f"  {t:<7} acc {c / tot:.4f}  soft acc {s / tot:.4f}  (n={tot})")
+    print("per workflow:")
+    for d, (c, tot, s) in sorted(by_domain.items()):
+        print(f"  {d:<28} acc {c / tot:.4f}  soft acc {s / tot:.4f}  (n={tot})")
+
+    if latency:
+        measure_latency(models[0], packers[0], rows, device)
+    return result
+
+
+def measure_latency(model, packer, rows, device, n_single=50, n_batch=200, batch_size=32):
+    """Latency protocol: p50 single-question and batched per-question throughput.
+    Model must already be on `device` and warmed up by at least a few calls."""
+    import time
+
+    qs = [(r["state"], q) for r in rows for q in r["questions"]]
+
+    def run_one(s, q):
+        pq = {"name": q["name"], "type": q["type"], "instructions": q.get("instructions", q["type"]), "options": q["options"], "answer": q["answer"]}
+        ids, an, _ = packer.pack(s, pq, model.cfg["max_len"])
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            model(torch.tensor([ids], device=device), torch.zeros(1, len(ids), dtype=torch.bool, device=device), torch.tensor([an], device=device))
+        if device == "cuda":
+            torch.cuda.synchronize()
+        return (time.perf_counter() - t0) * 1000
+
+    for s, q in qs[:3]:  # warmup
+        run_one(s, q)
+    times = sorted(run_one(s, q) for s, q in qs[:n_single])
+    print(f"\nlatency single-question p50: {times[len(times) // 2]:.1f} ms  (n={n_single}, warmup 3)")
+
+    # batched throughput: pack n_batch questions, pad to common length, one forward per batch
+    items = qs[:n_batch]
+    t0 = time.perf_counter()
+    done = 0
+    for i in range(0, len(items), batch_size):
+        chunk = items[i : i + batch_size]
+        packed = []
+        for s, q in chunk:
+            pq = {"name": q["name"], "type": q["type"], "instructions": q.get("instructions", q["type"]), "options": q["options"], "answer": q["answer"]}
+            packed.append(packer.pack(s, pq, model.cfg["max_len"]))
+        L = max(len(p[0]) for p in packed)
+        A = max(len(p[1]) for p in packed)
+        B = len(packed)
+        ids = torch.zeros(B, L, dtype=torch.long, device=device)
+        pmask = torch.ones(B, L, dtype=torch.bool, device=device)
+        apos = torch.zeros(B, A, dtype=torch.long, device=device)
+        for j, (pid, pan, _) in enumerate(packed):
+            ids[j, : len(pid)] = torch.tensor(pid, device=device)
+            pmask[j, : len(pid)] = False
+            apos[j, : len(pan)] = torch.tensor(pan, device=device)
+        with torch.no_grad():
+            model(ids, pmask, apos)
+        done += B
+    if device == "cuda":
+        torch.cuda.synchronize()
+    total_ms = (time.perf_counter() - t0) * 1000
+    print(f"latency batched: {total_ms / done:.1f} ms/question  (n={done}, batch={batch_size})")
+    print(f"throughput: {done / (total_ms / 1000):.0f} questions/sec")
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", default=None)
+    p.add_argument("--ckpts", default=None, help="comma-separated checkpoints (ensemble)")
+    p.add_argument("--data-dir", default="data/typed")
+    p.add_argument("--sharpen", type=float, default=1.0, help="confidence exponent gamma; >1 sharpens distributions")
+    p.add_argument("--latency", action="store_true", help="also measure single p50 and batched throughput")
+    args = p.parse_args()
+    ckpts = ([c.strip() for c in args.ckpts.split(",") if c.strip()]
+             if args.ckpts else [args.checkpoint])
+    evaluate_metrics(ckpts, args.data_dir, gamma=args.sharpen, latency=args.latency)

oev/calibrate.py

@@ -0,0 +1,22 @@
+import torch
+
+
+def apply_temperature(logits, temperature):
+    return logits / temperature
+
+
+def fit_temperature_from_logits(logits, labels, max_iter=200, lr=0.05):
+    logits = logits.detach().float().cpu()
+    labels = labels.detach().cpu()
+    log_t = torch.zeros(1, requires_grad=True)
+    opt = torch.optim.LBFGS([log_t], max_iter=max_iter, lr=lr)
+    labels = labels.long()
+
+    def closure():
+        opt.zero_grad()
+        loss = torch.nn.functional.cross_entropy(logits / log_t.exp(), labels)
+        loss.backward()
+        return loss
+
+    opt.step(closure)
+    return float(log_t.exp().item())

oev/convert.py

@@ -0,0 +1,61 @@
+import json
+from pathlib import Path
+
+AG_LABELS = ["world", "sports", "business", "sci/tech"]
+EMOTION_LABELS = ["sadness", "joy", "love", "anger", "fear", "surprise"]
+
+
+def convert_rows(rows, labels, domain, qname, instructions):
+    for i, r in enumerate(rows):
+        yield {
+            "id": f"{domain}-{i:06d}",
+            "domain": domain,
+            "state": r["text"],
+            "questions": [
+                {
+                    "name": qname,
+                    "type": "choice",
+                    "options": list(labels),
+                    "answer": labels[r["label"]],
+                    "instructions": instructions,
+                }
+            ],
+        }
+
+
+def download_ag_news():
+    from datasets import load_dataset
+
+    return load_dataset("fancyzhx/ag_news")
+
+
+def download_emotion():
+    from datasets import load_dataset
+
+    return load_dataset("dair-ai/emotion")
+
+
+def build_domain(hf_split_rows, labels, domain, qname, instructions):
+    return list(convert_rows(hf_split_rows, labels, domain, qname, instructions))
+
+
+def write_jsonl(rows, path):
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    with open(p, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+
+
+if __name__ == "__main__":
+    ag = download_ag_news()
+    em = download_emotion()
+    ag_train = build_domain(ag["train"], AG_LABELS, "ag_news", "topic", "Which topic does this news article belong to?")
+    write_jsonl(ag_train[:-2000], "data/ag_news/train.jsonl")
+    write_jsonl(ag_train[-2000:], "data/ag_news/valid.jsonl")
+    write_jsonl(build_domain(ag["test"], AG_LABELS, "ag_news", "topic", "Which topic does this news article belong to?"), "data/ag_news/test.jsonl")
+    em_train = build_domain(em["train"], EMOTION_LABELS, "emotion", "emotion", "Which emotion does this text express?")
+    write_jsonl(em_train[:-1000], "data/emotion/train.jsonl")
+    write_jsonl(em_train[-1000:], "data/emotion/valid.jsonl")
+    write_jsonl(build_domain(em["test"], EMOTION_LABELS, "emotion", "emotion", "Which emotion does this text express?"), "data/emotion/test.jsonl")
+    print("ag_news train/valid/test written; emotion train/valid/test written")

oev/convert_banking77.py

@@ -0,0 +1,88 @@
+"""Convert PolyAI/banking77 into OEV format.
+
+Banking77: 13,083 customer-service queries, 77 fine-grained intents.
+This is the high-cardinality test: 77 options in ONE choice question.
+Laya's published number is 0.425 (head token budget squeezes each label
+to 3-4 tokens); Jev publishes 0.870. Our anchor mechanism gives every
+option full tokens, so this is our most winnable remaining benchmark row.
+
+One choice question per state, all 77 intents as options with their
+human-readable intent names as option text.
+
+Usage:
+    python -m oev.convert_banking77
+Writes:
+    data/banking77/train.jsonl  (~10k cases, 1 question each)
+    data/banking77/valid.jsonl  (1,000 cases)
+    data/banking77/test.jsonl   (3,000 cases)
+"""
+
+import json
+from pathlib import Path
+
+
+def download_banking77():
+    from datasets import load_dataset
+
+    try:
+        ds = load_dataset("PolyAI/banking77")
+    except RuntimeError:
+        # new datasets versions dropped script support; use the auto-converted parquet branch
+        ds = load_dataset("PolyAI/banking77", revision="refs/convert/parquet")
+    return ds["train"], ds["test"]
+
+
+def intent_names():
+    """Official BANKING77 intent names in the dataset's label-id order,
+    fetched from the dataset's schema so it can never drift from the ids."""
+    import json
+    import urllib.request
+
+    url = "https://huggingface.co/datasets/PolyAI/banking77/resolve/main/dataset_infos.json"
+    data = json.load(urllib.request.urlopen(url, timeout=30))
+    names = list(data.values())[0]["features"]["label"]["names"]
+    assert len(names) == 77, f"expected 77 official intents, got {len(names)}"
+    return names
+
+
+def write_jsonl(rows, path):
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    with open(p, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+
+
+def convert(rows, names):
+    """One case per query: a single choice question with ALL 77 intents."""
+    n = len(names)
+    for r in rows:
+        label = int(r["label"])
+        if not (0 <= label < n):
+            continue
+        yield {
+            "id": f"b77-{r.get('id', label)}",
+            "domain": "banking77",
+            "state": r["text"],
+            "questions": [{
+                "name": "intent",
+                "type": "choice",
+                "instructions": "Which banking intent does this customer query belong to? Choose the single closest match.",
+                "options": names,
+                "answer": names[label],
+            }],
+        }
+
+
+if __name__ == "__main__":
+    train, test = download_banking77()
+    names = intent_names()
+    assert len(names) == 77, f"expected 77 intents, got {len(names)}"
+
+    train_rows = list(convert(train, names))
+    # carve a validation split off the train pool
+    write_jsonl(train_rows[:-1000], "data/banking77/train.jsonl")
+    write_jsonl(train_rows[-1000:], "data/banking77/valid.jsonl")
+    test_rows = list(convert(test, names))
+    write_jsonl(test_rows, "data/banking77/test.jsonl")
+    print(f"banking77 train/valid/test written: {len(train_rows) - 1000}/{1000}/{len(test_rows)} cases (77 options each)")

oev/convert_typed.py

@@ -0,0 +1,124 @@
+import json
+from pathlib import Path
+
+NOUL_OPTIONS = ["no", "yes"]
+
+
+def render_state(raw):
+    try:
+        obj = json.loads(raw) if isinstance(raw, str) else raw
+    except (json.JSONDecodeError, TypeError):
+        return str(raw)
+    if isinstance(obj, str):
+        return obj
+    if isinstance(obj, dict):
+        return "\n".join(f"{k}: {v}" for k, v in obj.items())
+    return str(obj)
+
+
+def score_levels(gold_q, criteria):
+    probs = gold_q.get("probabilities", {})
+    int_levels = [int(k) for k in probs if str(k).lstrip("-").isdigit()]
+    n = len(criteria) if isinstance(criteria, list) else 4
+    label = int(gold_q["label"])
+    return max(n, label + 1, (max(int_levels) + 1 if int_levels else 0))
+
+
+def soft_target(gold_q, options):
+    probs = gold_q.get("probabilities", {})
+    total = 0.0
+    target = []
+    for o in options:
+        p = probs.get(o)
+        if p is None and o == "no":
+            p = probs.get("false")
+        elif p is None and o == "yes":
+            p = probs.get("true")
+        p = float(p) if p is not None else 0.0
+        target.append(p)
+        total += p
+    if total <= 0.0:
+        return None
+    return [p / total for p in target]
+
+
+def _enrich(option, desc):
+    d = (desc or "").strip()
+    return f"{option}: {d}" if d else option
+
+
+def convert_question(qname, q, gold_q):
+    t = q["type"]
+    crit = q.get("criteria", {})
+    base = {"name": qname, "type": t, "instructions": q.get("instructions", t)}
+    if t == "choice":
+        keys = list(crit.keys())
+        if str(gold_q["label"]) not in keys:
+            return None
+        target = soft_target(gold_q, keys)
+        options = [_enrich(k, crit[k]) for k in keys]
+        answer = options[keys.index(str(gold_q["label"]))]
+    elif t == "noul":
+        options = NOUL_OPTIONS
+        target = soft_target(gold_q, options)
+        answer = "yes" if str(gold_q["label"]).lower() == "true" else "no"
+    elif t == "score":
+        n = score_levels(gold_q, crit)
+        raw = [str(i) for i in range(n)]
+        target = soft_target(gold_q, raw)
+        descs = crit if isinstance(crit, list) else []
+        options = [_enrich(str(i), descs[i] if i < len(descs) else "") for i in range(n)]
+        label = int(gold_q["label"])
+        if label >= n:
+            return None
+        answer = options[label]
+    else:
+        return None
+    cq = {**base, "options": options, "answer": answer}
+    if target is not None:
+        cq["target"] = target
+    return cq
+
+
+def convert_rows_typed(rows):
+    for row in rows:
+        questions = json.loads(row["questions"]) if isinstance(row["questions"], str) else row["questions"]
+        gold = json.loads(row["gold"]) if isinstance(row["gold"], str) else row["gold"]
+        qds = []
+        for qname, q in questions.items():
+            if qname not in gold:
+                continue
+            cq = convert_question(qname, q, gold[qname])
+            if cq is not None:
+                qds.append(cq)
+        if not qds:
+            continue
+        yield {
+            "id": f"typed-{row['id']}",
+            "domain": row.get("workflow", "typed"),
+            "state": render_state(row["state"]),
+            "questions": qds,
+        }
+
+
+def download_typed():
+    from datasets import load_dataset
+
+    return load_dataset("LocalLLaMA/typed-decisions", "all")
+
+
+def write_jsonl(rows, path):
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    with open(p, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+
+
+if __name__ == "__main__":
+    ds = download_typed()
+    train_rows = list(convert_rows_typed(ds["train"]))
+    write_jsonl(train_rows[:-200], "data/typed/train.jsonl")
+    write_jsonl(train_rows[-200:], "data/typed/valid.jsonl")
+    write_jsonl(list(convert_rows_typed(ds["test"])), "data/typed/test.jsonl")
+    print(f"typed train/valid/test written: {len(train_rows) - 200}/{200}/{len(ds['test'])} cases")