File size: 6,699 Bytes
e79f133 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """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:]))
|