| |
| """Stage 3b: canonical zone map for EVERY record of every tier. |
| |
| Single source of truth consumed by stages 4/5/6 (shard-level `zone` columns |
| are only preliminary). Rules: |
| - literary pristine: bucket from stage 3 |
| - papyri (ddbdp/dclp, ANY tier): TM digit rule, with chunk suffixes (#N) |
| stripped before the TM lookup [doublecheck D2 fix] |
| - inscriptions: PHI digit rule (from stage 1 zones) |
| - repaired literary: the bucket of its pristine sibling work/volume (same |
| work_prefix) when one exists, else TRAIN [doublecheck D3 fix: repaired |
| pages of a volume leave train exactly when the volume is val/test] |
| - bronze: TRAIN |
| |
| Output: work/zones_final.parquet (rid, zone int8) over all tiers. |
| """ |
| import glob |
| import json |
| import os |
| import sys |
| from collections import Counter, defaultdict |
|
|
| import pyarrow as pa |
| import pyarrow.parquet as pq |
| import xxhash |
| import orjson |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from common import ZONE_PTEST, ZONE_PVAL, ZONE_TRAIN, work_prefix |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| DDBDP_JSONL = os.path.expandvars("$STOICHEIA_DATA/clean/ddbdp.jsonl") |
| PAPYRI_TM_JSONL = os.path.expandvars("$STOICHEIA_DATA/data/papyri_clean.jsonl") |
|
|
|
|
| def digit_zone(numstr): |
| d = str(numstr).rstrip()[-1] |
| if d == "3": |
| return ZONE_PTEST |
| if d == "4": |
| return ZONE_PVAL |
| return ZONE_TRAIN |
|
|
|
|
| def build_tm_map(): |
| tm_by_base = {} |
| with open(PAPYRI_TM_JSONL, "rb") as f: |
| for line in f: |
| r = orjson.loads(line) |
| tm_by_base[r["file"]] = str(r["TM"]) |
| id2tm = {} |
| with open(DDBDP_JSONL, "rb") as f: |
| for line in f: |
| r = orjson.loads(line) |
| tm = tm_by_base.get(r["file"].rsplit("/", 1)[-1]) |
| if tm is not None: |
| id2tm[r["id"]] = tm |
| return id2tm |
|
|
|
|
| def main(): |
| id2tm = build_tm_map() |
|
|
| |
| lz = pq.read_table(os.path.join(ROOT, "work", "literary_zones.parquet"), |
| columns=["rid", "zone"]) |
| lit_zone = dict(zip(lz["rid"].to_pylist(), |
| [int(z) for z in lz["zone"].to_pylist()])) |
| cl = pq.read_table(os.path.join(ROOT, "work", "clusters.parquet"), |
| columns=["rid", "prefix"]) |
| prefix_bucket = {} |
| conflicts = 0 |
| for rid, pf in zip(cl["rid"].to_pylist(), cl["prefix"].to_pylist()): |
| b = lit_zone.get(rid) |
| if b is None: |
| continue |
| if pf in prefix_bucket and prefix_bucket[pf] != b: |
| conflicts += 1 |
| prefix_bucket[pf] = b |
|
|
| stats = Counter() |
| rids_all, zones_all = [], [] |
| seen = set() |
| dup_rids = 0 |
| for tier in ("pristine", "repaired", "bronze", "inscriptions"): |
| for p in sorted(glob.glob(os.path.join(ROOT, "work", "sentences", |
| tier, "shard_*.parquet"))): |
| t = pq.read_table(p, columns=["rid", "source", "zone"]) |
| for rid, src, z1 in zip(t["rid"].to_pylist(), |
| t["source"].to_pylist(), |
| t["zone"].to_pylist()): |
| base = rid.split("#")[0] |
| if src == "dclp": |
| z = digit_zone(base.split("_")[0]) |
| elif src == "ddbdp": |
| tm = id2tm.get(base) or id2tm.get(base.split("_")[0]) |
| if tm is None: |
| stats["ddbdp_tm_fallback_" + tier] += 1 |
| z = digit_zone(str(xxhash.xxh64_intdigest(base) % 10)) |
| else: |
| z = digit_zone(tm) |
| elif src == "phi": |
| z = int(z1) |
| elif tier == "pristine": |
| z = lit_zone.get(rid, ZONE_TRAIN) |
| if rid not in lit_zone: |
| stats["literary_unassigned"] += 1 |
| elif tier == "repaired": |
| z = prefix_bucket.get(work_prefix(src, rid), ZONE_TRAIN) |
| if z != ZONE_TRAIN: |
| stats["repaired_sibling_bucketed"] += 1 |
| else: |
| z = ZONE_TRAIN |
| if rid in seen: |
| dup_rids += 1 |
| seen.add(rid) |
| rids_all.append(rid) |
| zones_all.append(z) |
| stats["zone_%d_%s" % (z, tier)] += 1 |
|
|
| out = pa.table({"rid": rids_all, |
| "zone": pa.array(zones_all, type=pa.int8())}) |
| pq.write_table(out, os.path.join(ROOT, "work", "zones_final.parquet"), |
| compression="zstd") |
| summary = {"records": len(rids_all), "duplicate_rids_across_tiers": dup_rids, |
| "prefix_bucket_conflicts": conflicts, |
| "counts": dict(sorted(stats.items()))} |
| with open(os.path.join(ROOT, "work", "stage3b_stats.json"), "w") as f: |
| json.dump(summary, f, indent=2) |
| print(json.dumps(summary, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
