File size: 2,698 Bytes
7ed86c3 | 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 | """Shard-level region_id/century_id plumbing: written by insc build_shards.py, read by
ShardSet/MultiTierLoader, must survive _window()'s slicing (scalar, not per-char)."""
import sys
from pathlib import Path
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from train.data import DataConfig, MultiTierLoader, ShardSet, TierSpec, UNK_CENTURY, UNK_REGION
PLANES = ("chars", "boundary", "dia", "cap", "punct")
def _write_shard(d, n_records=6, rec_len=300, with_meta=True):
d.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(0)
bufs = {p: [] for p in PLANES}
rows = {"offset": [], "length": [], "tier": [], "clean": [], "dup_frac": []}
if with_meta:
rows["region_id"] = []
rows["century_id"] = []
off = 0
for i in range(n_records):
for p in PLANES:
bufs[p].append(rng.integers(0, 4, rec_len).astype(np.uint8))
rows["offset"].append(off); rows["length"].append(rec_len)
rows["tier"].append("iphi"); rows["clean"].append(1.0); rows["dup_frac"].append(0.0)
if with_meta:
rows["region_id"].append(i % 3)
rows["century_id"].append(i % 5)
off += rec_len
for p in PLANES:
np.concatenate(bufs[p]).tofile(d / f"{p}.bin")
pq.write_table(pa.table(rows), d / "index.parquet")
def test_shardset_reads_region_century_columns(tmp_path):
d = tmp_path / "shard_with_meta"
_write_shard(d, with_meta=True)
ss = ShardSet(str(d))
assert list(ss.region_id[:3]) == [0, 1, 2]
assert list(ss.century_id[:5]) == [0, 1, 2, 3, 4]
def test_shardset_defaults_unk_when_columns_absent(tmp_path):
"""Old shards (GCB pretraining, gold/silver/bronze) never had these columns."""
d = tmp_path / "shard_no_meta"
_write_shard(d, with_meta=False)
ss = ShardSet(str(d))
assert (ss.region_id == UNK_REGION).all()
assert (ss.century_id == UNK_CENTURY).all()
def test_multitier_loader_yields_region_century_and_survives_windowing(tmp_path):
d = tmp_path / "shard_windowed"
_write_shard(d, n_records=4, rec_len=5000, with_meta=True) # long enough to force _window
cfg = DataConfig(tiers={"iphi": TierSpec(str(d), 1.0, tier_filter="iphi")},
window_chars=256, seed=0)
loader = MultiTierLoader(cfg)
recs = list(loader.records(20))
assert recs
for r in recs:
assert len(r["chars"]) <= 256
assert isinstance(r["region_id"], (int, np.integer))
assert isinstance(r["century_id"], (int, np.integer))
assert r["region_id"] in (0, 1, 2)
assert r["century_id"] in (0, 1, 2, 3, 4)
|