| """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 <corpora_dir> <rows_dir> |
| <corpora_dir> 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 <package_dir> <rows_dir> <predictions.jsonl> |
| python bengu_public_eval.py score <predictions.jsonl> [<other_predictions.jsonl> ...] |
| 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:])) |
|
|
