File size: 29,110 Bytes
18d6af8
 
 
 
 
 
 
 
66e6f6e
 
 
 
18d6af8
66e6f6e
 
 
 
 
 
 
 
 
 
 
 
 
18d6af8
66e6f6e
18d6af8
 
7f57bf0
18d6af8
 
 
 
66e6f6e
18d6af8
 
66e6f6e
 
 
 
484f967
 
 
 
18d6af8
 
 
 
1f792cf
18d6af8
 
484f967
 
18d6af8
1f792cf
 
 
484f967
 
 
 
 
18d6af8
 
66e6f6e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7f57bf0
 
 
 
 
 
 
 
 
 
 
66e6f6e
18d6af8
66e6f6e
18d6af8
 
66e6f6e
18d6af8
 
66e6f6e
18d6af8
 
1f792cf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18d6af8
1f792cf
 
18d6af8
1f792cf
 
 
 
66e6f6e
1f792cf
 
 
 
 
 
66e6f6e
 
484f967
1f792cf
 
484f967
 
1f792cf
 
 
 
 
 
 
 
 
7f57bf0
 
18d6af8
1f792cf
 
 
 
 
 
 
484f967
1f792cf
 
 
 
 
484f967
1f792cf
 
 
484f967
1f792cf
 
 
 
 
 
484f967
 
66e6f6e
 
 
 
 
1f792cf
 
18d6af8
 
484f967
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7f57bf0
484f967
 
 
 
 
 
7f57bf0
 
484f967
 
 
 
 
 
7f57bf0
 
484f967
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18d6af8
 
 
484f967
 
 
 
 
 
 
66e6f6e
 
1b5ea42
 
 
 
 
18d6af8
 
 
 
1b5ea42
 
18d6af8
1b5ea42
 
 
 
1f792cf
1b5ea42
 
1f792cf
1b5ea42
 
 
 
 
 
 
1f792cf
1b5ea42
 
 
 
 
 
 
1f792cf
1b5ea42
 
 
 
 
 
 
1f792cf
1b5ea42
484f967
7f57bf0
 
 
484f967
 
18d6af8
66e6f6e
 
484f967
 
 
 
 
 
 
 
 
 
7f57bf0
 
 
 
 
 
 
484f967
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66e6f6e
484f967
 
98e4956
7f57bf0
 
484f967
 
7f57bf0
 
 
 
 
 
 
 
484f967
 
 
 
 
18d6af8
7f57bf0
484f967
 
 
 
 
 
 
 
 
7f57bf0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
484f967
 
 
 
7f57bf0
 
484f967
 
 
 
 
 
7f57bf0
 
484f967
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18d6af8
 
66e6f6e
 
1f792cf
484f967
66e6f6e
 
484f967
98e4956
18d6af8
484f967
18d6af8
484f967
 
7f57bf0
 
 
 
 
484f967
66e6f6e
484f967
 
 
 
 
 
 
 
66e6f6e
484f967
 
 
 
 
7f57bf0
 
 
 
 
18d6af8
484f967
66e6f6e
484f967
 
 
 
18d6af8
66e6f6e
18d6af8
66e6f6e
 
 
98e4956
18d6af8
 
 
 
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
# Phase 2 contamination audit (§3.3): mechanical 13-token overlap between the mix and the eval tasks.
#
# What §3.3 permits, precisely: "benchmark validation/dev material may be used mechanically: a script
# computes overlap counts and you never read or learn from the items themselves. Test splits stay
# untouched until Phase 6." This script obeys that literally, and it is written so that violating it is
# hard: the only values that ever leave the process are counts and rates per source. No item text, no
# lengths, no samples, nothing a human could read as a question. The print surface is a dict of integers.
#
# Which splits are the reference, and why it is broader than the minimum: for most of these tasks the
# harness draws its few-shot exemplars from TRAIN (ARC, HellaSwag, PIQA, Winogrande) or from MMLU's `dev`,
# so overlap with those is contamination by construction. The reference set is {train, validation, dev}
# where present, per task, and NEVER `test`.
#
# Why this file does not use Python sets of hashed grams (it did, and it was wrong): 1.23 B tokens in
# 13-token windows is ~1.2e9 windows, and `hashlib.sha256` per window in CPython is ~1-2 us each, so the
# audit would have taken hours to days and simply never finished inside a CPU session. Two changes make it
# tractable without weakening it:
#   1. Gram fingerprints are a vectorised multiplicative rolling hash (mod 2^64) computed per document with
#      numpy, not a Python loop. Cost is ~20 uint64 ops per token, in C.
#   2. The reference is ONE sorted uint64 array plus a uint16 bitmask giving the tasks each gram belongs to,
#      queried with np.searchsorted. Memory is ~1 GB, not the ~7 GB a Python set of 90M ints needs, and
#      per-reference sets are unnecessary.
# Collision behaviour is stated rather than hidden: 64-bit fingerprints over ~1.2e9 mix grams and ~1e8
# reference grams give an expected spurious-match count of mix_grams x ref_grams / 2^64, reported as
# `expected_false_positives`. A raw count below that is noise, not evidence of contamination; above it,
# the offending source is investigated.
#
# Read-only over the mix. CPU session. No credentials in this file.

import argparse
import hashlib
import json
import os
import struct

import numpy as np

K = 13
P = np.uint64(0x100000001B3)          # FNV-1a prime, odd, so invertible mod 2^64
PINV = np.uint64(pow(int(P), -1, 1 << 64))
MODMASK = np.uint64((1 << 64) - 1)

# task -> (repo, config(s) to try in order, splits treated as reference material; 'test' deliberately
# absent). Split names were read off the dataset server on 2026-09-19 rather than assumed: gsm8k has
# train/test only, truthfulqa's repo is `truthfulqa/truthful_qa` and exposes `validation` only, and
# `sylinrl/TruthfulQA` -- the name the first version of this file used -- no longer exists.
REFS = [
    ("arc_easy", "allenai/ai2_arc", ["ARC-Easy"], ["train", "validation"]),
    ("arc_challenge", "allenai/ai2_arc", ["ARC-Challenge"], ["train", "validation"]),
    ("hellaswag", "Rowan/hellaswag", [None], ["train", "validation"]),
    ("piqa", "baber/piqa", [None], ["train", "validation"]),
    ("winogrande", "allenai/winogrande", ["winogrande_xl"], ["train", "validation"]),
    ("mmlu", "cais/mmlu", ["all"], ["dev", "validation"]),
    ("truthfulqa", "truthfulqa/truthful_qa", ["multiple_choice", "generation"], ["validation"]),
    ("gsm8k", "openai/gsm8k", ["main"], ["train"]),
]
TEXT_KEYS = ("question", "text", "choices", "story", "sentence", "prompt", "ctx", "input", "query",
             "answer", "solution", "reasoning", "explanation", "goal", "sol1", "sol2", "abstractions",
             "nonabstractions", "mc1_targets", "mc2_targets")
# A reference split that yields fewer rows than this did not load, it merely failed to raise: the first
# run of this audit reported `piqa::default::train: 0 items -> 0 grams` with no exception, which a
# coverage-blind verdict would have scored as "no contamination found in PIQA".
MIN_ITEMS_PER_SPLIT = 200



def window_hashes(ids, k=K):
    """uint64 fingerprint of every k-length window of `ids` (uint16 array), vectorised.

    h_i = sum_{j=i}^{i+k-1} x_j * P^(i+k-1-j)  mod 2^64. Via the prefix recurrence
    C_i = C_{i-1}*P + x_i  and the identity h_i = C_{i+k-1} - C_{i-1} * P^k. C is computed as
    P^i * cumsum(x_j * PINV^j), which is what makes it a few numpy passes instead of a Python loop.
    """
    n = int(ids.shape[0])
    if n < k:
        return np.empty(0, dtype=np.uint64)
    x = ids.astype(np.uint64)

    def powers(base, m):
        # [base^0, base^1, ..., base^(m-1)] mod 2^64. np.cumprod of a constant vector would start at
        # base^1, which shifts every window hash and silently breaks the match against the reference.
        out = np.empty(m, dtype=np.uint64)
        out[0] = np.uint64(1)
        if m > 1:
            out[1:] = np.cumprod(np.full(m - 1, base, dtype=np.uint64))
        return out

    pow_p, pow_pi = powers(P, n), powers(PINV, n)
    s = np.cumsum(x * pow_pi, dtype=np.uint64)
    c = s * pow_p                                    # c[i] = C_i, prefix hash ending at i
    pk = np.uint64(pow(int(P), k, 1 << 64))
    prev = np.concatenate((np.zeros(1, dtype=np.uint64), (c[:n - k] * pk)))
    h = (c[k - 1:] - prev[:n - k + 1]) & MODMASK
    return h


def brief_error(e):
    """Short, publishable description of a loader failure.

    This report is shipped with the mix, and an exception string from a dataset reader can quote a cell of
    benchmark text -- exactly the thing §3.3 keeps out of every artifact. Quoted spans are replaced.
    """
    import re
    msg = re.sub(r"'[^']*'|\"[^\"]*\"", "'…'", str(e).replace("\n", " "))
    return f"{type(e).__name__}: {msg[:140]}"


def flatten(v, out):
    if isinstance(v, str):
        out.append(v)
    elif isinstance(v, dict):
        for x in v.values():
            flatten(x, out)
    elif isinstance(v, (list, tuple)):
        for x in v:
            flatten(x, out)


SERVER = "https://datasets-server.huggingface.co"
_SIZE_CACHE = {}


def server_sizes(repo, cfg):
    """Row counts the dataset server reports for each split of one config, or {} if it cannot be asked.

    This exists because a streaming reader can stop early without raising, and because a loader that yields
    0 rows of a 16,113-row split (exactly what `baber/piqa` did under the Kaggle image's `datasets`) is
    indistinguishable from a clean result unless the expected count is checked. Never trust the count from
    memory: the first version of this file expected 14,042 MMLU validation rows, and 14,042 is MMLU's
    *test* split -- the one §3.3 forbids reading. `validation` is 1,531.
    """
    import urllib.parse
    import urllib.request
    ck = (repo, cfg)
    if ck in _SIZE_CACHE:
        return _SIZE_CACHE[ck]
    q = {"dataset": repo}
    if cfg:
        q["config"] = cfg
    out = {}
    try:
        with urllib.request.urlopen(SERVER + "/size?" + urllib.parse.urlencode(q), timeout=90) as r:
            d = json.load(r)
        out = {s["split"]: int(s.get("num_rows", 0)) for s in d.get("size", {}).get("splits", [])}
    except Exception as e:
        print(f"  sizes API {repo}/{cfg or ''}: {type(e).__name__}", flush=True)
    _SIZE_CACHE[ck] = out
    return out


def api_rows(repo, cfg, sp, cap):
    """Page the public rows endpoint (100 per request, full=true so no cell is silently truncated)."""
    import urllib.parse
    import urllib.request
    off = 0
    while off < cap:
        q = {"dataset": repo, "split": sp, "offset": off, "limit": 100, "full": "true"}
        if cfg:
            q["config"] = cfg
        try:
            with urllib.request.urlopen(SERVER + "/rows?" + urllib.parse.urlencode(q),
                                        timeout=120) as r:
                d = json.load(r)
        except Exception as e:
            print(f"  rows API {repo}/{cfg or ''}/{sp} @{off}: {type(e).__name__}", flush=True)
            return
        batch = d.get("rows") or []
        if not batch:
            return
        for item in batch:
            yield item.get("row", item)
        off += len(batch)


def ingest(tk, rows, cap):
    """Encode each row's text fields into gram windows. Returns (items kept, list of gram arrays)."""
    g, n = [], 0
    for row in rows:
        if n >= cap:
            break
        parts = []
        for k in TEXT_KEYS:
            if k in row:
                flatten(row[k], parts)
        txt = " ".join(parts)
        if not txt.strip():
            continue
        n += 1
        ids = np.asarray(tk.encode(txt, add_special_tokens=False).ids, dtype=np.uint16)
        h = window_hashes(ids)
        if h.size:
            g.append(h)
    return n, g


def build_reference(tk, max_items_per_split):
    """Returns (sorted uint64 reference grams, parallel uint16 task bitmask, per-split item counts,
    errors, per-split expected row counts). Only counts and hashes are ever held or printed.

    A split is accepted only if the rows actually read reach 90 % of what the dataset server says the
    split holds; short reads are retried through the HTTP rows endpoint and, if still short, recorded as
    loader problems -- which fail MEASURED rather than scoring as "no contamination found".
    """
    grams, masks = [], []
    counts, errors, srv = {}, {}, {}
    try:
        import datasets
    except Exception as e:
        datasets = None
        errors["import datasets"] = f"{type(e).__name__}: {str(e)[:120]}"
    for ti, (task, repo, configs, splits) in enumerate(REFS):
        bit = np.uint16(1 << ti)
        for sp in splits:
            key = f"{task}::{repo}::{sp}"
            best = None
            for cfg in configs:
                ckey = f"{key}[{cfg or 'default'}]"
                want = server_sizes(repo, cfg).get(sp)
                srv[key] = want
                n, g = 0, []
                if datasets is not None:
                    try:
                        ds = (datasets.load_dataset(repo, cfg, split=sp, streaming=True) if cfg
                              else datasets.load_dataset(repo, split=sp, streaming=True))
                        n, g = ingest(tk, ds, max_items_per_split)
                    except Exception as e:
                        errors[ckey] = brief_error(e)
                        print(f"  ref {ckey}: UNREADABLE {brief_error(e)}", flush=True)
                        continue
                if want and n < int(0.9 * want):
                    print(f"  ref {ckey}: {n} of {want} rows via datasets -> rows-API fallback",
                          flush=True)
                    n2, g2 = ingest(tk, api_rows(repo, cfg, sp, max_items_per_split),
                                    max_items_per_split)
                    if n2 > n:
                        n, g = n2, g2
                counts[ckey] = n
                print(f"  ref {ckey}: {n} items of {want} -> {sum(a.size for a in g):,} grams",
                      flush=True)
                if best is None or n > best[1]:
                    best = (cfg, n, g, want)
                if n >= MIN_ITEMS_PER_SPLIT and (not want or n >= int(0.9 * want)):
                    break
            if best is None or best[1] < MIN_ITEMS_PER_SPLIT:
                errors[key] = (f"no config of {repo} yielded >= {MIN_ITEMS_PER_SPLIT} rows for {sp} "
                               f"(best {(best[1] if best else 0)} of {srv.get(key)})")
                continue
            cfg, n, g, want = best
            if want and n < int(0.9 * want):
                errors[key] = f"covered {n} of {want} rows in {sp} -- under-read, not calling it clean"
                continue
            if g:
                arr = np.concatenate(g)
                grams.append(arr)
                masks.append(np.full(arr.shape, bit, dtype=np.uint16))
    if not grams:
        raise SystemExit("no reference grams built at all -- the audit would be a pass over nothing")
    allg = np.concatenate(grams)
    allm = np.concatenate(masks)
    order = np.argsort(allg, kind="stable")
    return allg[order], allm[order], counts, errors, srv




def gram_lookup(g, ref, refmask):
    """(number of hits, the task bitmask of the matched reference grams) for one document.

    numpy only, because this runs ~1.1 M times. `searchsorted` into the sorted reference table plus the
    parallel bitmask is what makes per-task attribution possible without holding eight sets.
    """
    if g.size == 0:
        return 0, None
    pos = np.clip(np.searchsorted(ref, g), 0, ref.size - 1)
    eq = ref[pos] == g
    if not eq.any():
        return 0, None
    return int(eq.sum()), refmask[pos[eq]]


def per_task_counts(bits):
    return [int((bits & np.uint16(1 << ti)).astype(np.uint16).sum()) for ti in range(len(REFS))]


def scan_docs(ref, refmask, docs, seen=None):
    """Fold an iterator of (ordinal_within_source, raw_bytes) into counts.

    The ordinal is the only identity a document ever gets here: a filter pass can drop it by index later
    without anything in this process knowing what it said. `hit_docs_over_half` is the column that decides
    interpretation -- contamination by verbatim inclusion shows up as documents whose grams are *mostly*
    matched, while one shared boilerplate sentence shows up as a single gram in a handful of documents.
    Pass one `seen` set across a whole mix, otherwise the duplicate count is per-shard and only a lower
    bound (found in review).
    """
    import time
    import hashlib
    t0 = time.time()
    s = {"docs": 0, "tokens": 0, "grams": 0, "hit_docs": 0, "hits": 0, "hit_docs_over_half": 0,
         "hit_tokens": 0, "dup_docs": 0, "dup_tokens": 0, "per_task": [0] * len(REFS), "ordinals": []}
    if seen is None:
        seen = set()
    names = [t[0] for t in REFS]
    for idx, raw in docs:
        b = bytes(raw)
        ids = np.frombuffer(b, dtype=np.uint16)
        s["docs"] += 1
        s["tokens"] += int(ids.size)
        dg = hashlib.blake2b(b, digest_size=8).digest()
        if dg in seen:                      # counted, never removed: de-duplication is out of scope (D-014)
            s["dup_docs"] += 1
            s["dup_tokens"] += int(ids.size)
        else:
            seen.add(dg)
        g = window_hashes(ids)
        s["grams"] += int(g.size)
        h, bits = gram_lookup(g, ref, refmask)
        if h:
            s["hit_docs"] += 1
            s["hits"] += h
            s["hit_tokens"] += int(ids.size)
            s["per_task"] = [a + c for a, c in zip(s["per_task"], per_task_counts(bits))]
            s["ordinals"].append(idx)
            if h * 2 >= int(g.size):
                s["hit_docs_over_half"] += 1
    s["seconds"] = round(time.time() - t0, 1)
    s["per_task_named"] = {n: c for n, c in zip(names, s["per_task"]) if c}
    return s


def doc_stream(BM, shard_paths):
    """(ordinal, raw document bytes) over a fixed list of shards. Ordinal order is the shard order in
    the source's own record.json, which is exactly what the merge replays -- so an index computed here
    still means the same document after a session restart."""
    n = 0
    for p in shard_paths:
        for d in BM.iter_shard_docs(p):
            yield n, d
            n += 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="/kaggle/working/mixroot")
    ap.add_argument("--stage-subdir", default="stage",
                    help="per-source staged shards live here under --root; scanning them is what makes "
                         "attribution possible, because every merged shard contains all sources at once")
    ap.add_argument("--max-items-per-split", type=int, default=50000)
    ap.add_argument("--write-filter", default="",
                    help="directory for <source>.u8 exclusion masks + filter.json (read back by "
                         "build_mix --exclude-dir)")
    ap.add_argument("--out", default="", help="write the report JSON here (publish_mix expects "
                                             "mixroot/audit.json)")
    ap.add_argument("--ref-cache", default="",
                    help="directory to keep the built reference gram table in. Streaming the eight "
                         "reference sets is the slow half of this script, and a re-audit after a filter "
                         "pass rebuilds exactly the same table -- so it is reused only if the task list, "
                         "item cap, k and tokenizer sha all still match.")
    a = ap.parse_args()

    from tokenizers import Tokenizer
    from huggingface_hub import hf_hub_download
    tjp = hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json")
    tk = Tokenizer.from_file(tjp)
    tk.no_truncation()
    import hashlib
    tk_sha = hashlib.sha256(open(tjp, "rb").read()).hexdigest()

    want_meta = {"k": K, "max_items": a.max_items_per_split, "tokenizer_sha256": tk_sha,
                 "text_keys": list(TEXT_KEYS),
                 "refs": [[t[0], t[1], list(t[2]), list(t[3])] for t in REFS]}
    ref = refmask = None
    refcounts = errors = srv = {}
    if a.ref_cache and os.path.exists(os.path.join(a.ref_cache, "refmeta.json")):
        try:
            if json.load(open(os.path.join(a.ref_cache, "refmeta.json"))) == want_meta:
                ref = np.load(os.path.join(a.ref_cache, "ref.npy"))
                refmask = np.load(os.path.join(a.ref_cache, "refmask.npy"))
                refcounts = json.load(open(os.path.join(a.ref_cache, "refcounts.json")))
                errors = json.load(open(os.path.join(a.ref_cache, "errors.json")))
                srv = json.load(open(os.path.join(a.ref_cache, "srv.json")))
                print(f"reference: loaded from cache, {ref.size:,} grams", flush=True)
        except Exception as e:
            print(f"reference cache unusable ({type(e).__name__}: {str(e)[:120]}); rebuilding",
                  flush=True)
            ref = None
    if ref is None:
        print("building reference gram table (only counts ever leave this process)", flush=True)
        ref, refmask, refcounts, errors, srv = build_reference(tk, a.max_items_per_split)
        if a.ref_cache:
            os.makedirs(a.ref_cache, exist_ok=True)
            np.save(os.path.join(a.ref_cache, "ref.npy"), ref)
            np.save(os.path.join(a.ref_cache, "refmask.npy"), refmask)
            json.dump(want_meta, open(os.path.join(a.ref_cache, "refmeta.json"), "w"))
            json.dump(refcounts, open(os.path.join(a.ref_cache, "refcounts.json"), "w"))
            json.dump(errors, open(os.path.join(a.ref_cache, "errors.json"), "w"))
            json.dump(srv, open(os.path.join(a.ref_cache, "srv.json"), "w"))
            print(f"reference cached under {a.ref_cache}", flush=True)
    tasks_covered = sorted({k.split("::")[0] for k, v in refcounts.items()
                            if v >= MIN_ITEMS_PER_SPLIT} -
                           {k.split("::")[0] for k in errors
                            if "::" in k and "[" not in k and not k.startswith("stage::")})
    print(f"reference: {ref.size:,} grams, {len(tasks_covered)}/8 tasks covered, "
          f"{len(errors)} loader problems", flush=True)

    import build_mix as BM

    # ---- per-source attribution, over the staged (pre-merge) shards ------------------------------
    # The first version attributed merged hits as `c // len(rec["sources"])` to every source named in the
    # shard. All 141 shards name all 15 sources, so every source received an almost identical number and
    # the table looked like a finding while being arithmetically meaningless.
    stage_root = os.path.join(a.root, a.stage_subdir)
    by_source, masks = {}, {}
    if os.path.isdir(stage_root):
        for key in sorted(os.listdir(stage_root)):
            sd = os.path.join(stage_root, key)
            rd = os.path.join(sd, "record.json")
            if not os.path.isdir(sd):
                continue
            if not os.path.exists(rd):
                # Absent record means the source never finished staging or was built without a Hub
                # repo. Treating it as "nothing to scan" would merge it unfiltered (review finding).
                errors[f"stage::{key}"] = "no record.json -- source is incomplete or unpublished"
                print(f"  stage {key}: NO record.json, source left unscanned", flush=True)
                continue
            rec = json.load(open(rd))
            paths = [os.path.join(a.root, s["file"]) for s in rec.get("shards", [])]
            absent = [p for p in paths if not os.path.exists(p)]
            if absent:
                errors[f"stage::{key}"] = f"{len(absent)} staged shard(s) missing on disk"
                print(f"  stage {key}: {len(absent)} shard(s) missing, skipped", flush=True)
                continue
            s = scan_docs(ref, refmask, doc_stream(BM, paths))
            s["shards"] = len(paths)
            s["stage_complete"] = bool(rec.get("complete"))
            by_source[key] = s
            masks[key] = (np.asarray(sorted(s["ordinals"]), dtype=np.int64), s["docs"])
            print(f"  stage {key}: {s['docs']:,} docs {s['tokens']:,} tok -> hits={s['hits']:,} in "
                  f"{s['hit_docs']:,} docs ({s['hit_docs_over_half']:,} of them >50 % matched) "
                  f"[{s['seconds']}s] {json.dumps(s['per_task_named'])}", flush=True)
            s.pop("ordinals")
    else:
        print(f"no staged sources under {stage_root}: attribution unavailable, totals only", flush=True)

    # ---- the verdict is always measured on the bytes that will actually train ---------------------
    man_path = os.path.join(a.root, "manifest.json")
    hits = {}
    docs = grams_seen = mix_tokens = hit_docs = hit_docs_over_half = 0
    dup_docs = dup_tokens = 0
    missing = []
    val = {"docs": 0, "tokens": 0, "grams": 0, "hit_docs": 0, "hits": 0, "hit_docs_over_half": 0}
    manifest_digest = None
    if os.path.exists(man_path):
        man = json.load(open(man_path))
        # Bound this report to the exact bytes it scanned: publish_mix re-computes the digest from the
        # manifest it is about to ship, so an audit.json carried over from a different mix -- even one of
        # the same total size -- is refused instead of certifying it.
        manifest_digest = hashlib.sha256(json.dumps(
            [[s["file"], s.get("tokens"), s.get("docs")] for s in man["shards"]] +
            [["VAL:" + s["file"], s.get("tokens"), s.get("docs")] for s in man.get("val_shards", [])],
            sort_keys=True).encode()).hexdigest()
        seen_all = set()      # one set across the whole mix, so duplicates are counted, not floored
        for rec in man["shards"]:
            p = os.path.join(a.root, rec["file"])
            if not os.path.exists(p):
                # A shard not on disk is unmeasured, not clean (E-017's shape).
                missing.append(rec["file"])
                continue
            s = scan_docs(ref, refmask, doc_stream(BM, [p]), seen=seen_all)
            docs += s["docs"]
            mix_tokens += s["tokens"]
            grams_seen += s["grams"]
            hit_docs += s["hit_docs"]
            hit_docs_over_half += s["hit_docs_over_half"]
            dup_docs += s["dup_docs"]
            dup_tokens += s["dup_tokens"]
            for n, c in s["per_task_named"].items():
                hits[n] = hits.get(n, 0) + c
        # The held-out split is reported on, so it must be scanned too. A benchmark document sitting in
        # validation would flatter every validation-PPL number this project publishes, and nothing else
        # in the pipeline reads those bytes (review finding).
        for rec in man.get("val_shards", []):
            p = os.path.join(a.root, rec["file"])
            if not os.path.exists(p):
                missing.append(rec["file"])
                continue
            v = scan_docs(ref, refmask, doc_stream(BM, [p]), seen=seen_all)
            for k2 in ("docs", "tokens", "grams", "hit_docs", "hits", "hit_docs_over_half"):
                val[k2] += int(v.get(k2, 0))
        total_overlap = sum(hits.values()) + val["hits"]
        print(f"  mix: {docs:,} docs, {grams_seen:,} grams, {sum(hits.values()):,} train hits in "
              f"{hit_docs:,} docs ({hit_docs_over_half:,} majority-hit); val: {val['docs']:,} docs, "
              f"{val['hits']:,} hits in {val['hit_docs']:,} docs", flush=True)
    else:
        total_overlap = 0
        print(f"no manifest at {a.root}: merged mix not scanned", flush=True)

    grams_all = grams_seen + val["grams"]
    expected_fp = grams_all * int(np.unique(ref).size) / float(1 << 64)
    noise_allow = max(10.0, 10.0 * expected_fp)
    per_task_named = {t[0]: hits.get(t[0], 0) for t in REFS}

    if a.write_filter:
        os.makedirs(a.write_filter, exist_ok=True)
        summary = {"k": K, "criterion": "any 13-gram of the document appears in a reference "
                                        "train/validation/dev split",
                   "scanned_sources": sorted(masks), "sources": {}}
        for key, (ord_sorted, n_docs) in masks.items():
            m = np.zeros(n_docs, dtype=np.uint8)
            if ord_sorted.size:
                m[ord_sorted] = 1
            m.tofile(os.path.join(a.write_filter, key + ".u8"))
            s = by_source[key]
            summary["sources"][key] = {"docs": s["docs"], "tokens": s["tokens"],
                                       "dropped_docs": int(ord_sorted.size),
                                       "dropped_tokens": int(s["hit_tokens"]),
                                       "kept_docs": int(n_docs - ord_sorted.size)}
        with open(os.path.join(a.write_filter, "filter.json"), "w") as f:
            json.dump(summary, f, indent=1, default=str)
        print(f"filter masks written to {a.write_filter} for {len(masks)} source(s)", flush=True)

    contaminated_sources = {k: v["hit_docs"] / max(1, v["docs"]) for k, v in by_source.items()
                            if v["hit_docs"]}
    report = {
        "k": K,
        "fingerprint": "multiplicative rolling hash mod 2^64, P=0x100000001B3",
        "reference_items_by_split": refcounts,
        "reference_rows_expected_by_split": srv,
        "reference_loader_problems": errors,
        "reference_grams_distinct": int(np.unique(ref).size),
        "reference_grams_total": int(ref.size),
        "tasks_covered": tasks_covered,
        "shards_missing": missing,
        "mix_documents_audited": docs,
        "mix_tokens_audited": mix_tokens,
        "mix_grams_audited": grams_seen,
        "mix_documents_with_any_hit": hit_docs,
        "mix_documents_majority_hit": hit_docs_over_half,
        "val_documents_audited": val["docs"],
        "val_tokens_audited": val["tokens"],
        "val_documents_with_any_hit": val["hit_docs"],
        "val_hits": val["hits"],
        "mix_bytes_sha256": manifest_digest,
        "overlap_by_task": per_task_named,
        "overlap_total": total_overlap,
        "by_source": by_source,
        "source_hit_document_fraction": {k: round(v, 6) for k, v in sorted(
            contaminated_sources.items(), key=lambda kv: -kv[1])},
        "duplicate_documents": dup_docs,
        "duplicate_tokens": dup_tokens,
        "expected_false_positives": round(expected_fp, 3),
        "noise_allowance": round(noise_allow, 3),
        "overlap_exceeds_noise": bool(total_overlap > noise_allow),
        "note": ("counts only; no item text read or printed. Reference splits are train/validation/dev "
                 "per task -- never test, per §3.3. overlap_by_task sums gram hits, so one verbatim "
                 "benchmark document contributes as many hits as it has windows; read "
                 "mix_documents_majority_hit and source_hit_document_fraction to judge severity. "
                 "Duplicate counting is diagnostic only: de-duplication is out of scope (D-014), "
                 "contamination removal is not (§3.3)."),
        "COVERED_ALL_SHARDS": bool(not missing and docs > 0 and val["docs"] > 0),
        "FILTER_MASKS_WRITTEN": sorted(masks) if a.write_filter else None,
        "MEASURED": bool(docs > 0 and grams_seen > 0 and val["docs"] > 0 and not errors
                         and len(tasks_covered) == 8 and manifest_digest
                         and (not a.write_filter or masks)),
    }
    report["AUDIT_PASSED"] = bool(report["MEASURED"] and total_overlap <= noise_allow)
    txt = json.dumps(report, indent=1, default=str)
    print(f"VERDICT AUDIT_PASSED={report['AUDIT_PASSED']} MEASURED={report['MEASURED']} "
          f"docs={docs:,} overlap={total_overlap} hit_docs={hit_docs:,} "
          f"majority_hit_docs={hit_docs_over_half:,} expected_fp={expected_fp:.3f} "
          f"tasks={len(tasks_covered)}/8 problems={len(errors)}")
    print("AUDIT_JSON_BEGIN")
    print(txt)
    print("AUDIT_JSON_END")
    if a.out:
        with open(a.out, "w") as f:
            f.write(txt)
    raise SystemExit(0 if report["AUDIT_PASSED"] else (6 if report["MEASURED"] else 7))


if __name__ == "__main__":
    main()