File size: 2,350 Bytes
e6f80f3 | 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 | """python -m fastfill.tools.ablation_check <base data dir> <constrained data dir>
The constrained build must differ from the base build only by the `constraints` key of some train rows:
same complete rows (including source, flags and all message metadata) after removing only that key from
the constrained user JSON; dev/test files byte-identical (one evaluation serves both)."""
import hashlib
from itertools import zip_longest
import json
import sys
def _normalized(row, remove_constraints=False):
"""Keep the entire row; normalize only the intended user JSON payload."""
user = json.loads(row["messages"][1]["content"])
if remove_constraints:
user = {k: v for k, v in user.items() if k != "constraints"}
normalized = {**row, "messages": [
{**m, "content": user} if i == 1 else m for i, m in enumerate(row["messages"])]}
# JSON booleans and numbers must not compare equal (Python's True == 1 does).
return json.dumps(normalized, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
def _sha256(path):
digest = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1024 * 1024), b""):
digest.update(chunk)
return digest.digest()
def main(base, cons):
diff = n_rows = n_cons = 0
with open(f"{base}/train.jsonl", encoding="utf-8") as a, open(f"{cons}/train.jsonl", encoding="utf-8") as b:
for lx, ly in zip_longest(a, b):
assert lx is not None and ly is not None, "train row counts differ"
x, y = json.loads(lx), json.loads(ly)
n_rows += 1
n_cons += "constraints" in json.loads(y["messages"][1]["content"])
diff += _normalized(x) != _normalized(y, remove_constraints=True)
same_eval = all(_sha256(f"{base}/{f}") == _sha256(f"{cons}/{f}")
for f in ("dev.jsonl", "test.jsonl", "dev_rooms.jsonl", "test_rooms.jsonl",
"dev_constrained_rooms.jsonl", "test_constrained_rooms.jsonl"))
print(json.dumps({"train_rows": n_rows, "rows_differing_beyond_constraints": diff,
"constrained_rows_in_" + cons.rstrip("/").split("/")[-1]: n_cons, "eval_files_identical": same_eval}))
assert diff == 0 and same_eval, "not a clean ablation"
if __name__ == "__main__":
main(sys.argv[1], sys.argv[2])
|