File size: 4,571 Bytes
5952424 | 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 | #!/usr/bin/env python3
"""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 # train-always records never appear in val/test
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)
# merge: concat -> sort -> OR-reduce duplicate keys
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()
|