Bosun vs slimemold — head-to-head on slimemold's own edge_quality benchmark (DialAM-2024 / QT30, same 5 episodes, seed 42). Isolates the edge-judgment step.

bosun_slimemold_edge_eval_modal.py

#!/usr/bin/env -S uv run --with modal
"""Head-to-head: Bosun vs slimemold's edge-mapping on slimemold's OWN edge_quality benchmark
(DialAM-2024 / QT30). We replicate his episode selection (seed 42, identical 5 episodes / 29 gold
relations), then give Bosun the GOLD claim nodes and have it judge every ordered pair for
supports / contradicts. This isolates the EDGE-judgment step (what Bosun would replace) and removes
the claim-extraction bottleneck that caps his reported edge recall (52%, limited by 76% claim recall).

slimemold's reported numbers on these episodes: edge recall 52% (15/29), relation-type acc 100%,
edge precision 10% (he flags precision as misleading on all-pairs).

Run:  cd repos/atlas-warrant-service
      mise exec -- uv run --with modal modal run --detach ../../bosun_slimemold_edge_eval_modal.py
"""
import modal

app = modal.App("bosun-slimemold-edge-eval")
vol = modal.Volume.from_name("atlas-warrant")
image = (modal.Image.debian_slim(python_version="3.11")
         .pip_install("torch>=2.2.0", "transformers>=4.51.0", "peft>=0.11.0", "numpy>=1.24"))

MODELS = {
    "4B_v11": dict(base="Qwen/Qwen3-Reranker-4B", adapter="/ckpt/lora_adapters_v11_4b",
                   tok="/ckpt/tokenizer_v11_4b", cfg="/ckpt/serving_v11_4b.json"),
    "SmolLM135": dict(base="HuggingFaceTB/SmolLM2-135M", adapter="/ckpt/lora_adapters_smollm135",
                      tok="/ckpt/tokenizer_smollm135", cfg="/ckpt/serving_smollm135.json"),
    "XS_v11": dict(base="Qwen/Qwen3-Reranker-0.6B", adapter="/ckpt/lora_adapters_v11",
                   tok="/ckpt/tokenizer_v11", cfg="/ckpt/serving_v11.json"),
    "XS_v10": dict(base="Qwen/Qwen3-Reranker-0.6B", adapter="/ckpt/lora_adapters_v10",
                   tok="/ckpt/tokenizer_v10", cfg="/ckpt/serving_v10.json"),
    "4B_v10_4b": dict(base="Qwen/Qwen3-Reranker-4B", adapter="/ckpt/lora_adapters_v10_4b",
                      tok="/ckpt/tokenizer_v10_4b", cfg="/ckpt/serving_v10_4b.json"),
}

QUERY = "These two findings share the specified relationship."
SUPPORT = ("Connected ONLY if FINDING A states a premise, reason, or evidence that SUPPORTS or "
           "justifies the claim in FINDING B. Otherwise not connected.")
CONTRADICT = ("Connected ONLY if FINDING A states something that CONTRADICTS or conflicts with the "
              "claim in FINDING B. Otherwise not connected.")

# --- slimemold's episode selection, replicated verbatim (seed 42) so we hit the SAME 5 episodes ---
DATASET_URL = "http://dialam.arg.tech/res/files/dataset.zip"
MIN_I_NODES, MAX_I_NODES, MIN_RA, MIN_CA, MAX_LOCUTIONS, NUM_EPISODES, SEED = 8, 35, 2, 1, 50, 5, 42


def load_episodes():
    import os, json, zipfile, urllib.request, random
    ddir = "/tmp/dialam-dataset/dataset"
    if not os.path.isdir(ddir):
        urllib.request.urlretrieve(DATASET_URL, "/tmp/d.zip")
        with zipfile.ZipFile("/tmp/d.zip") as z:
            z.extractall("/tmp/dialam-dataset")

    def parse_episode(path):
        data = json.load(open(path))
        nodes = {n["nodeID"]: n for n in data["nodes"]}
        from_edges, to_edges = {}, {}
        for e in data["edges"]:
            from_edges.setdefault(e["fromID"], []).append(e["toID"])
            to_edges.setdefault(e["toID"], []).append(e["fromID"])
        i_nodes = {n["nodeID"]: n["text"] for n in data["nodes"] if n["type"] == "I"}
        gold = []
        for s_type, rel in [("RA", "supports"), ("CA", "contradicts")]:
            for s in [n for n in data["nodes"] if n["type"] == s_type]:
                sid = s["nodeID"]
                srcs = [n for n in to_edges.get(sid, []) if nodes.get(n, {}).get("type") == "I"]
                tgts = [n for n in from_edges.get(sid, []) if nodes.get(n, {}).get("type") == "I"]
                for a in srcs:
                    for b in tgts:
                        gold.append({"from_id": a, "to_id": b, "relation": rel})
        return {"i_nodes": i_nodes, "gold_relations": gold}

    cands = []
    for fn in os.listdir(ddir):
        if not fn.endswith(".json"):
            continue
        data = json.load(open(os.path.join(ddir, fn)))
        ns = data["nodes"]
        ic = sum(1 for n in ns if n["type"] == "I")
        rac = sum(1 for n in ns if n["type"] == "RA")
        cac = sum(1 for n in ns if n["type"] == "CA")
        lc = sum(1 for n in ns if n["type"] == "L")
        if MIN_I_NODES <= ic <= MAX_I_NODES and rac >= MIN_RA and cac >= MIN_CA and lc <= MAX_LOCUTIONS:
            cands.append(os.path.join(ddir, fn))
    random.seed(SEED)
    chosen = random.sample(cands, min(NUM_EPISODES, len(cands)))
    return [parse_episode(p) for p in chosen]


@app.function(image=image, gpu="L40S", volumes={"/ckpt": vol}, timeout=2400)
def run():
    import json, torch, numpy as np
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from peft import PeftModel

    episodes = load_episodes()
    n_gold = sum(len(e["gold_relations"]) for e in episodes)
    n_claims = sum(len(e["i_nodes"]) for e in episodes)
    print(f"episodes={len(episodes)}  gold_claims={n_claims}  gold_relations={n_gold}", flush=True)

    def score_with(m):
        cfg = json.loads(open(m["cfg"]).read())
        tok = AutoTokenizer.from_pretrained(m["tok"], padding_side="left")
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        base = AutoModelForCausalLM.from_pretrained(cfg["base_model"], trust_remote_code=True,
                                                    attn_implementation="sdpa", torch_dtype=torch.bfloat16)
        model = PeftModel.from_pretrained(base, m["adapter"]).merge_and_unload().eval().cuda()
        pids = tok(cfg["prefix"], add_special_tokens=False)["input_ids"]
        sids = tok(cfg["suffix"], add_special_tokens=False)["input_ids"]
        budget = cfg["max_len"] - len(pids) - len(sids)

        def score(pairs, instruct):
            items = []
            for i, (a, b) in enumerate(pairs):
                body = f"<Instruct>: {instruct}\n<Query>: {QUERY}\n<Document>: FINDING A:\n{a}\n\nFINDING B:\n{b}"
                cids = tok(body, add_special_tokens=False)["input_ids"][:budget]
                items.append((i, pids + cids + sids))
            items.sort(key=lambda x: len(x[1]))
            out = [0.0] * len(pairs); k = 0; cap = 24000
            while k < len(items):
                bs = max(1, cap // len(items[k][1])); chunk = items[k:k + bs]
                mx = max(len(x[1]) for x in chunk); pad = tok.pad_token_id
                bb = [[pad] * (mx - len(x[1])) + x[1] for x in chunk]
                am = [[0] * (mx - len(x[1])) + [1] * len(x[1]) for x in chunk]
                with torch.no_grad():
                    ii = torch.tensor(bb, dtype=torch.long, device="cuda")
                    aa = torch.tensor(am, dtype=torch.long, device="cuda")
                    lg = model(input_ids=ii, attention_mask=aa, logits_to_keep=1).logits[:, -1, :].float()
                    s = 1.0 / (1.0 + np.exp(-(lg[:, cfg["yes_id"]] - lg[:, cfg["no_id"]]).cpu().numpy()))
                for (o, _), sc in zip(chunk, s):
                    out[o] = float(sc)
                k += len(chunk)
            return out

        # gather all ordered pairs across episodes, with global gold sets
        all_pairs, meta = [], []          # meta: (episode_idx, from_id, to_id)
        for ei, ep in enumerate(episodes):
            ids = list(ep["i_nodes"].keys())
            for a in ids:
                for b in ids:
                    if a != b:
                        all_pairs.append((ep["i_nodes"][a], ep["i_nodes"][b]))
                        meta.append((ei, a, b))
        sup = score(all_pairs, SUPPORT)
        con = score(all_pairs, CONTRADICT)
        del model, base; torch.cuda.empty_cache()
        return cfg, meta, sup, con

    # gold lookups (direction-agnostic, like slimemold's matcher; type normalized support/conflict)
    def norm(rel):
        return "support" if rel == "supports" else "conflict"
    gold_pairs = {}          # (ei, frozenset{a,b}) -> normalized type
    gold_directed = set()    # (ei, a, b, type) for the strict directional view
    for ei, ep in enumerate(episodes):
        for g in ep["gold_relations"]:
            gold_pairs[(ei, frozenset((g["from_id"], g["to_id"])))] = norm(g["relation"])
            gold_directed.add((ei, g["from_id"], g["to_id"], norm(g["relation"])))

    report = {"episodes": len(episodes), "gold_claims": n_claims, "gold_relations": n_gold,
              "slimemold_ref": {"edge_recall": 0.52, "type_acc": 1.00, "edge_precision": 0.10}, "models": {}}

    for tag, m in MODELS.items():
        cfg, meta, sup, con = score_with(m)
        # per unordered pair: best support/contradict over both directions
        upair = {}   # (ei, frozenset) -> {"sup":max, "con":max}
        sdir = {}    # (ei,a,b)->sup score for directional support check
        for (ei, a, b), su, co in zip(meta, sup, con):
            key = (ei, frozenset((a, b)))
            d = upair.setdefault(key, {"sup": 0.0, "con": 0.0})
            d["sup"] = max(d["sup"], su); d["con"] = max(d["con"], co)
            sdir[(ei, a, b)] = su

        def metrics(thr):
            # predicted edges among ALL unordered pairs
            pred = {k: ("support" if v["sup"] >= v["con"] else "conflict")
                    for k, v in upair.items() if max(v["sup"], v["con"]) >= thr}
            # recall + type acc over gold
            found = sum(1 for k in gold_pairs if k in pred)
            type_ok = sum(1 for k, t in gold_pairs.items() if k in pred and pred[k] == t)
            recall = found / len(gold_pairs)
            type_acc = (type_ok / found) if found else 0.0
            precision = (sum(1 for k in pred if k in gold_pairs) / len(pred)) if pred else 0.0
            # strict directional support recall (RA only): A supports B in the gold direction
            ra = [(ei, a, b) for (ei, a, b, t) in gold_directed if t == "support"]
            dir_found = sum(1 for (ei, a, b) in ra if sdir.get((ei, a, b), 0) >= thr)
            dir_recall = dir_found / len(ra) if ra else 0.0
            return dict(thr=thr, recall=round(recall, 3), type_acc=round(type_acc, 3),
                        precision=round(precision, 3), n_pred=len(pred), dir_support_recall=round(dir_recall, 3))

        # AUROC: gold-edge vs non-edge over unordered pairs, score = max(sup,con) (threshold-free quality)
        ys, ss = [], []
        for k, v in upair.items():
            ys.append(1 if k in gold_pairs else 0); ss.append(max(v["sup"], v["con"]))
        pos = [s for s, y in zip(ss, ys) if y]; neg = [s for s, y in zip(ss, ys) if not y]
        if pos and neg:
            w = sum((p > nn) + 0.5 * (p == nn) for p in pos for nn in neg)
            auroc = w / (len(pos) * len(neg))
        else:
            auroc = float("nan")
        # best-F1 threshold sweep
        sweep = [metrics(t) for t in [float(round(x, 2)) for x in np.arange(0.3, 0.96, 0.05)]]
        bestf1 = max(sweep, key=lambda r: (2 * r["recall"] * r["precision"] / (r["recall"] + r["precision"]))
                     if (r["recall"] + r["precision"]) else 0)
        report["models"][tag] = {"auroc_edge_vs_none": round(auroc, 3),
                                 "at_0.5": metrics(0.5), "best_f1": bestf1}
        r5 = report["models"][tag]["at_0.5"]
        print(f"[{tag}] AUROC(edge vs none)={auroc:.3f}  @0.5 recall={r5['recall']} type_acc={r5['type_acc']} "
              f"precision={r5['precision']}  dir_support_recall={r5['dir_support_recall']}", flush=True)

    import os
    os.makedirs("/ckpt/eval_out", exist_ok=True)
    open("/ckpt/eval_out/bosun_slimemold_edge.json", "w").write(json.dumps(report, indent=2))
    vol.commit()
    print("wrote /ckpt/eval_out/bosun_slimemold_edge.json", flush=True)
    return report


@app.local_entrypoint()
def main():
    import json
    print(json.dumps(run.remote(), indent=2))