File size: 2,531 Bytes
032a961 | 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 | #!/usr/bin/env python3
"""Probe: readout-vocabulary coverage curve (Zipf knee) to pick vocab size.
Replicates the exact EPT fit/calibrate/segment pipeline of SPIRA.build on
train.bin, counts patch frequencies, then measures on held-out streams
(test.bin, oneshot.bin) the fraction of target patches covered by the
top-V vocabulary for V in a grid. Writes JSON to data/vocab_probe.json.
"""
import json
import sys
import time
from collections import Counter
sys.path.insert(0, "/home/z/my-project/download/spira")
import numpy as np
from spira.ept import EPT
from spira.hashing import fnv1a64
BASE = "/home/z/my-project"
t0 = time.perf_counter()
train = open(f"{BASE}/data/train.bin", "rb").read()
test = open(f"{BASE}/data/test.bin", "rb").read()
oneshot = open(f"{BASE}/data/oneshot.bin", "rb").read()
ept = EPT(budget_bits=16.0, min_len=2, max_len=32, alpha=0.5)
ept.fit(train[:2_000_000])
budget = ept.calibrate(train[2_000_000:2_400_000], target_len=4.5)
print(f"EPT budget={budget:.2f} bits ({time.perf_counter()-t0:.0f}s)")
def ids_of(data: bytes) -> np.ndarray:
patches = ept.segment(data)
return np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches))
ids_train = ids_of(train)
ids_test = ids_of(test)
ids_one = ids_of(oneshot)
print(f"patches: train={len(ids_train)} test={len(ids_test)} oneshot={len(ids_one)} "
f"({time.perf_counter()-t0:.0f}s)")
uniq, cnt = np.unique(ids_train, return_counts=True)
# same deterministic tie-break as SPIRA.build: count desc, larger id first
order = np.lexsort((uniq, -cnt))
sorted_ids = uniq[order]
GRID = [4096, 6144, 8192, 12288, 16384, 24576, 32768]
rows = []
for V in GRID:
vocab = set(sorted_ids[:V].tolist())
cov_test = float(np.mean([int(i) in vocab for i in ids_test.tolist()]))
cov_one = float(np.mean([int(i) in vocab for i in ids_one.tolist()]))
mass_train = float(cnt[order][:V].sum() / cnt.sum())
rows.append({"vocab": V, "cov_test": round(cov_test, 4),
"cov_oneshot": round(cov_one, 4),
"train_mass": round(mass_train, 4)})
print(f"V={V:>6} held-out test coverage={cov_test:.4f} "
f"oneshot coverage={cov_one:.4f} train mass={mass_train:.4f}")
out = {"budget_bits": float(budget), "distinct_train": int(len(uniq)),
"n_train": int(len(ids_train)), "curve": rows,
"elapsed_s": round(time.perf_counter() - t0, 1)}
with open(f"{BASE}/data/vocab_probe.json", "w") as f:
json.dump(out, f, indent=1)
print("saved -> data/vocab_probe.json")
|