| |
| """Stage 4: global contamination index. |
| |
| For every record that can ever be val/test (all pristine records in buckets |
| 0-9 or PTEST/PVAL, plus Inscriptions_2 variants in PTEST/PVAL) emit hash keys: |
| - exact skeleton hash of each sentence |
| - bag hash (sorted words) of each sentence [catches reorderings] |
| - every word-8-gram hash over the record's whole word stream |
| (cross-sentence, so differing punctuation/splits can't hide a quote) |
| each mapped to the record's zone bit. Keys are merged by bitwise OR. |
| |
| Output: work/index_keys.npy (sorted uint64), work/index_masks.npy (uint16) |
| """ |
| import glob |
| import json |
| import os |
| import sys |
| from concurrent.futures import ProcessPoolExecutor |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from common import h64, NGRAM, ZONE_TRAIN |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| PARTS = os.path.join(ROOT, "work", "index_parts") |
|
|
|
|
| def load_literary_zones(): |
| t = pq.read_table(os.path.join(ROOT, "work", "literary_zones.parquet"), |
| columns=["rid", "zone"]) |
| return dict(zip(t["rid"].to_pylist(), t["zone"].to_pylist())) |
|
|
|
|
| def process_shard(args): |
| tier, path, out_path = args |
| zf = ZONES_CACHE.get("z") |
| t = pq.read_table(path, columns=["rid", "zone", "skels"]) |
| keys, masks = [], [] |
| n_contrib = 0 |
| for rid, _z1, skels in zip(t["rid"].to_pylist(), t["zone"].to_pylist(), |
| t["skels"].to_pylist()): |
| zone = zf.get(rid, -1) |
| if zone < 0 or zone >= ZONE_TRAIN: |
| continue |
| bit = 1 << zone |
| n_contrib += 1 |
| words_all = [] |
| for sk in skels: |
| w = sk.split() |
| keys.append(h64(sk)) |
| masks.append(bit) |
| keys.append(h64(" ".join(sorted(w)))) |
| masks.append(bit) |
| words_all.extend(w) |
| if len(words_all) >= NGRAM: |
| for i in range(len(words_all) - NGRAM + 1): |
| keys.append(h64(" ".join(words_all[i:i + NGRAM]))) |
| masks.append(bit) |
| k = np.array(keys, dtype=np.uint64) |
| m = np.array(masks, dtype=np.uint16) |
| np.savez(out_path, k=k, m=m) |
| return len(k), n_contrib |
|
|
|
|
| ZONES_CACHE = {} |
|
|
|
|
| def init_worker(zpath): |
| import pyarrow.parquet as pq2 |
| t = pq2.read_table(zpath, columns=["rid", "zone"]) |
| ZONES_CACHE["z"] = dict(zip(t["rid"].to_pylist(), t["zone"].to_pylist())) |
|
|
|
|
| def main(): |
| os.makedirs(PARTS, exist_ok=True) |
| zpath = os.path.join(ROOT, "work", "zones_final.parquet") |
| tasks = [] |
| for tier in ("pristine", "inscriptions"): |
| for p in sorted(glob.glob(os.path.join(ROOT, "work", "sentences", |
| tier, "shard_*.parquet"))): |
| out = os.path.join(PARTS, "%s_%s.npz" % |
| (tier, os.path.basename(p).split(".")[0])) |
| tasks.append((tier, p, out)) |
| print("%d shards" % len(tasks), flush=True) |
|
|
| workers = max(4, os.cpu_count() - 8) |
| total_keys = 0 |
| total_contrib = 0 |
| with ProcessPoolExecutor(max_workers=workers, initializer=init_worker, |
| initargs=(zpath,)) as ex: |
| for nk, nc in ex.map(process_shard, tasks, chunksize=1): |
| total_keys += nk |
| total_contrib += nc |
| print("raw keys: %d from %d contributing records" % |
| (total_keys, total_contrib), flush=True) |
|
|
| |
| parts = sorted(glob.glob(os.path.join(PARTS, "*.npz"))) |
| ks, ms = [], [] |
| for p in parts: |
| z = np.load(p) |
| ks.append(z["k"]) |
| ms.append(z["m"]) |
| k = np.concatenate(ks) |
| m = np.concatenate(ms) |
| del ks, ms |
| order = np.argsort(k, kind="stable") |
| k = k[order] |
| m = m[order] |
| del order |
| boundary = np.empty(len(k), dtype=bool) |
| boundary[0] = True |
| np.not_equal(k[1:], k[:-1], out=boundary[1:]) |
| starts = np.flatnonzero(boundary) |
| uk = k[starts] |
| um = np.bitwise_or.reduceat(m, starts) |
| np.save(os.path.join(ROOT, "work", "index_keys.npy"), uk) |
| np.save(os.path.join(ROOT, "work", "index_masks.npy"), um) |
| stats = {"raw_keys": int(total_keys), "unique_keys": int(len(uk)), |
| "contributing_records": int(total_contrib)} |
| with open(os.path.join(ROOT, "work", "stage4_stats.json"), "w") as f: |
| json.dump(stats, f, indent=2) |
| print(json.dumps(stats, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
