"""Bengu-1 public triage track: rebuild the 1,587 test rows from the public corpora and score predictions. Standard library only, except `predict`, which needs bengu_infer.py's dependencies. python bengu_public_eval.py extract holds the source files under their upstream names (public-track-rows.json, "sources"): 20030228_easy_ham.tar.bz2, 20030228_easy_ham_2.tar.bz2, 20030228_hard_ham.tar.bz2, 20030228_spam.tar.bz2, 20050311_spam_2.tar.bz2, phishing-2024, phishing-2025, enron_with_categories.tar.gz. Each file's sha256 is checked, and each test message's raw sha256 and input-view sha256 are checked against the row list. Writes one .eml per row. python bengu_public_eval.py predict python bengu_public_eval.py score [ ...] A predictions line: {"row_id", "category", "protected", "needs_reply"}. A missing row counts as a failure and is reported. Scores the full track and the clean subset ("clean": true, 1,473 rows). Metrics: group macro-F1 over the coarse groups present in the gold rows (a predicted category counts as its group), junk recall/precision, protected and needs_reply recall/precision where the dataset label defines them. """ import hashlib import json import mailbox import sys import tarfile from pathlib import Path HERE = Path(__file__).resolve().parent ROWS = HERE / "public-track-rows.json" def load_rows(): return json.loads(ROWS.read_text()) def _sha256_file(path): h = hashlib.sha256() with open(path, "rb") as f: for block in iter(lambda: f.read(1 << 20), b""): h.update(block) return h.hexdigest() def _file_name(row_id): return row_id.replace("/", "__").replace("#", "_") + ".eml" def _messages(corpora, sources): """(row_id, raw bytes) of every message in the pinned source files.""" for name, src in sources.items(): path = Path(corpora) / name if _sha256_file(path) != src["sha256"]: raise SystemExit(f"{name}: sha256 differs from the pinned file; refused") if src["kind"] == "tar": with tarfile.open(path) as t: for m in t: if not m.isfile(): continue rid = src["row_id_of_member"] parts = m.name.split("/") if rid == "spamassassin" and len(parts) == 2: yield f"spamassassin/{parts[0]}/{parts[1]}", t.extractfile(m).read() elif rid == "enron-berkeley" and len(parts) == 3 and parts[2].endswith(".txt"): yield f"enron-berkeley/{parts[1]}/{parts[2][:-4]}", t.extractfile(m).read() elif src["kind"] == "mbox": box = mailbox.mbox(str(path), create=False) for i, key in enumerate(box.iterkeys()): yield f"nazario/{src['row_id_of_member']}#{i}", box.get_bytes(key) def extract(corpora, out): sys.path.insert(0, str(HERE.parent / "inference")) from bengu_view import view spec = load_rows() want = {r["row_id"]: r for r in spec["rows"]} out = Path(out) out.mkdir(parents=True, exist_ok=True) found, bad_raw, view_ok = 0, 0, 0 for rid, raw in _messages(corpora, spec["sources"]): r = want.get(rid) if r is None: continue if hashlib.sha256(raw).hexdigest() != r["raw_sha256"]: bad_raw += 1 continue (out / _file_name(rid)).write_bytes(raw) found += 1 view_ok += hashlib.sha256(view(raw).encode("utf-8")).hexdigest() == r["view_sha256"] print(json.dumps({"rows": len(want), "extracted": found, "raw_sha256_mismatch": bad_raw, "input_view_sha256_match": view_ok, "missing": len(want) - found - bad_raw})) return 0 if found == len(want) else 1 def predict(package, rows_dir, dst): sys.path.insert(0, str(HERE.parent / "inference")) from bengu_infer import Bengu1 from bengu_view import view model = Bengu1(package) rows = load_rows()["rows"] with open(dst, "w") as fh: for i in range(0, len(rows), 32): chunk = rows[i:i + 32] preds = model.predict([view((Path(rows_dir) / _file_name(r["row_id"])).read_bytes()) for r in chunk]) for r, p in zip(chunk, preds): fh.write(json.dumps({"row_id": r["row_id"], **p}) + "\n") return 0 def _prf(gold, pred): tp = sum(g and p for g, p in zip(gold, pred)) fp = sum((not g) and p for g, p in zip(gold, pred)) fn = sum(g and not p for g, p in zip(gold, pred)) return {"recall": round(tp / (tp + fn), 4) if tp + fn else None, "precision": round(tp / (tp + fp), 4) if tp + fp else None, "positives": tp + fn} def metrics(rows, preds, groups): group_of = {c: g for g, cats in groups.items() for c in cats} pairs = [(r["gold"], preds[r["row_id"]]) for r in rows if r["row_id"] in preds] out = {"n": len(rows), "failed": len(rows) - len(pairs)} gp = [(g["group"], group_of.get(p["category"], "unknown")) for g, p in pairs if g["group"] is not None] f1s = {} for grp in groups: tp = sum(a == grp and b == grp for a, b in gp) fp = sum(a != grp and b == grp for a, b in gp) fn = sum(a == grp and b != grp for a, b in gp) if tp + fn: f1s[grp] = round(2 * tp / (2 * tp + fp + fn), 4) if tp else 0.0 out["group_macro_f1"] = round(sum(f1s.values()) / len(f1s), 4) if f1s else None out["group_f1"] = f1s out["junk"] = _prf([g["junk"] for g, _ in pairs], [group_of.get(p["category"]) == "junk" for _, p in pairs]) for key in ("protected", "needs_reply"): known = [(g[key], bool(p[key])) for g, p in pairs if g[key] is not None] out[key] = _prf([g for g, _ in known], [p for _, p in known]) return out def score(*paths): spec = load_rows() rows = spec["rows"] clean = [r for r in rows if r["clean"]] for path in paths: preds = {} for line in Path(path).read_text().splitlines(): if line.strip(): p = json.loads(line) preds[p["row_id"]] = p print(json.dumps({"predictions": str(path), "full": metrics(rows, preds, spec["groups"]), "clean": metrics(clean, preds, spec["groups"])})) return 0 if __name__ == "__main__": cmds = {"extract": extract, "predict": predict, "score": score} if len(sys.argv) < 2 or sys.argv[1] not in cmds: raise SystemExit(__doc__) sys.exit(cmds[sys.argv[1]](*sys.argv[2:]))