data/split_pipeline/06_materialize.py

#!/usr/bin/env python3
"""Stage 6: materialize all 10 folds to disk (single pass over the corpus).

Fold k:
  test  = bucket k (pristine) + PTEST (papyri via TM, inscriptions via PHI)
  val   = bucket (k+1)%10 + PVAL, cleaned against fold-k test zones
  train = everything else, with contaminated sentences excised and maximal
          runs of consecutive clean sentences re-stitched into segments

Test records are emitted verbatim. Val records are excised of any sentence
colliding with the fold's test zones (belt and braces for cluster misses).
Train records are excised of any sentence colliding with the fold's val+test
zones. Segments < MIN_SEG_CHARS (train: 100, val: 25) are dropped. Each
maximal run becomes its OWN record (no false word adjacencies across cuts).

Inscriptions: train-zone variants are emitted per-variant like normal text
records. PTEST/PVAL inscriptions are emitted as the full original row MINUS
the synthetic fields (synthetics of val/test PHI numbers appear nowhere).
"""
import glob
import json
import os
import sys
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor

import orjson
import pyarrow.parquet as pq
import zstandard as zstd

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import (fold_conflict_mask, N_BUCKETS, MASK_PTEST, MASK_PVAL,
                    ZONE_PTEST, ZONE_PVAL, ZONE_TRAIN, MIN_SEG_CHARS)

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUTDIR = os.path.expandvars(os.environ.get("FOLD_OUTDIR", "$STOICHEIA_DATA"))
MIN_VAL_CHARS = 25
INSCR_JSONL = os.path.join(ROOT, "raw", "Inscriptions_2",
                           "synthetic_editions_with_ithaca_text_fix.jsonl")
INSCR_REAL = ["edition", "with_diacritics", "without_diacritics", "ithaca_text"]
INSCR_SYNTH = {"synthetic", "synthetic_2"}


def test_mask(k):
    return (1 << k) | MASK_PTEST


def segments(text, starts, ends, masks, conflict, min_chars):
    """Yield (piece_text, was_cut) after excising sentences hitting conflict."""
    if not starts:
        return [], 0, 0
    keep = [not (m & conflict) for m in masks]
    if all(keep):
        return [text], 0, 0
    pieces = []
    cut_chars = 0
    i = 0
    n = len(starts)
    while i < n:
        if not keep[i]:
            cut_chars += ends[i] - starts[i]
            i += 1
            continue
        j = i
        while j + 1 < n and keep[j + 1]:
            j += 1
        piece = text[starts[i]:ends[j]].strip()
        if len(piece) >= min_chars:
            pieces.append(piece)
        else:
            cut_chars += len(piece)
        i = j + 1
    return pieces, cut_chars, 1


class FoldWriters:
    def __init__(self, base, tag):
        self.base = base
        self.tag = tag
        self.w = {}

    def get(self, fold, split):
        key = (fold, split)
        if key not in self.w:
            d = os.path.join(self.base, "fold_%d" % fold, "parts")
            os.makedirs(d, exist_ok=True)
            f = open(os.path.join(d, "%s-%s.jsonl.zst" % (split, self.tag)), "wb")
            cctx = zstd.ZstdCompressor(level=6)
            self.w[key] = (f, cctx.stream_writer(f))
        return self.w[key][1]

    def close(self):
        for f, s in self.w.values():
            s.close()
            f.close()


def emit(writer, obj):
    writer.write(orjson.dumps(obj) + b"\n")


def process_shard(args):
    tier, spath, mpath = args
    tag = tier + "-" + os.path.basename(spath).split(".")[0]
    st = pq.read_table(spath, columns=["rid", "source", "text", "starts", "ends"])
    mt = pq.read_table(mpath, columns=["rid", "zone", "masks"])
    assert st["rid"].to_pylist() == mt["rid"].to_pylist(), "shard misalignment"
    W = FoldWriters(OUTDIR, tag)
    stats = defaultdict(lambda: [0, 0, 0, 0])  # (fold,split,tier)->recs,chars,cut,dropped

    rids = st["rid"].to_pylist()
    sources = st["source"].to_pylist()
    texts = st["text"].to_pylist()
    starts_c = st["starts"].to_pylist()
    ends_c = st["ends"].to_pylist()
    zones = mt["zone"].to_pylist()
    masks_c = mt["masks"].to_pylist()

    for rid, source, text, starts, ends, zone, masks in zip(
            rids, sources, texts, starts_c, ends_c, zones, masks_c):
        inscr = tier == "inscriptions"
        for k in range(N_BUCKETS):
            if zone == k or zone == ZONE_PTEST:
                role = "test"
            elif zone == (k + 1) % N_BUCKETS or zone == ZONE_PVAL:
                role = "val"
            else:
                role = "train"
            if inscr and role != "train":
                continue  # val/test inscriptions are emitted centrally
            if role == "test":
                if tier != "pristine":
                    continue  # only pristine (incl. papyri) feeds test
                s = stats[(k, "test", tier)]
                emit(W.get(k, "test"), {"id": rid, "tier": tier,
                                        "source": source, "text": text})
                s[0] += 1
                s[1] += len(text)
                continue
            if role == "val" and tier != "pristine":
                continue
            conflict = test_mask(k) if role == "val" else fold_conflict_mask(k)
            min_chars = MIN_VAL_CHARS if role == "val" else MIN_SEG_CHARS
            pieces, cut, was_cut = segments(text, starts, ends, masks,
                                            conflict, min_chars)
            s = stats[(k, role, tier)]
            if not pieces:
                s[3] += 1
                s[2] += cut
                continue
            wr = W.get(k, role)
            if len(pieces) == 1 and not was_cut:
                emit(wr, {"id": rid, "tier": tier, "source": source,
                          "text": pieces[0]})
            else:
                for pi, piece in enumerate(pieces):
                    emit(wr, {"id": "%s#seg%d" % (rid, pi), "tier": tier,
                              "source": source, "text": piece})
            s[0] += len(pieces)
            s[1] += sum(len(p) for p in pieces)
            s[2] += cut
    W.close()
    return {"%d|%s|%s" % key: v for key, v in stats.items()}


def inscriptions_valtest():
    """Emit PTEST/PVAL inscriptions centrally: original row minus synthetics."""
    # per-PHI OR of all real-variant sentence masks (for val-vs-test cleaning)
    phi_mask = {}
    phi_zone = {}
    for p in sorted(glob.glob(os.path.join(ROOT, "work", "masks",
                                           "inscriptions", "*.parquet"))):
        t = pq.read_table(p, columns=["rid", "zone", "masks"])
        for rid, zone, masks in zip(t["rid"].to_pylist(), t["zone"].to_pylist(),
                                    t["masks"].to_pylist()):
            phi, field = rid.split(":", 1)
            if zone not in (ZONE_PTEST, ZONE_PVAL) or field in INSCR_SYNTH:
                continue
            m = 0
            for x in masks:
                m |= x
            phi_mask[phi] = phi_mask.get(phi, 0) | m
            phi_zone[phi] = zone

    W = FoldWriters(OUTDIR, "inscr-central")
    stats = defaultdict(lambda: [0, 0, 0, 0])
    with open(INSCR_JSONL, "rb") as f:
        for line in f:
            try:
                r = orjson.loads(line)
            except orjson.JSONDecodeError:
                r = json.loads(line)  # pandas-written rows with bare NaN
            phi = "phi%s" % r["PHI_ID"]
            zone = phi_zone.get(phi)
            if zone is None:
                continue
            out = {"id": phi, "tier": "inscriptions", "source": "phi"}
            for fld, v in r.items():
                if fld not in INSCR_SYNTH:
                    out[fld] = v
            if zone == ZONE_PTEST:
                for k in range(N_BUCKETS):
                    emit(W.get(k, "test"), out)
                    s = stats[(k, "test", "inscriptions")]
                    s[0] += 1
            else:  # PVAL: drop from fold k's val if it collides with test zones
                m = phi_mask.get(phi, 0)
                for k in range(N_BUCKETS):
                    if m & test_mask(k):
                        stats[(k, "val", "inscriptions")][3] += 1
                        continue
                    emit(W.get(k, "val"), out)
                    stats[(k, "val", "inscriptions")][0] += 1
    W.close()
    return {"%d|%s|%s" % key: v for key, v in stats.items()}


def main():
    os.makedirs(OUTDIR, exist_ok=True)
    tasks = []
    for tier in ("pristine", "repaired", "bronze", "inscriptions"):
        for spath in sorted(glob.glob(os.path.join(
                ROOT, "work", "sentences", tier, "shard_*.parquet"))):
            mpath = os.path.join(ROOT, "work", "masks", tier,
                                 os.path.basename(spath))
            tasks.append((tier, spath, mpath))
    print("%d shards" % len(tasks), flush=True)

    agg = defaultdict(lambda: [0, 0, 0, 0])
    workers = max(4, os.cpu_count() - 8)
    with ProcessPoolExecutor(max_workers=workers) as ex:
        futs = [ex.submit(process_shard, t) for t in tasks]
        futs.append(ex.submit(inscriptions_valtest))
        for i, f in enumerate(futs):
            for key, v in f.result().items():
                a = agg[key]
                for j in range(4):
                    a[j] += v[j]
            if (i + 1) % 50 == 0:
                print("  %d/%d" % (i + 1, len(futs)), flush=True)

    with open(os.path.join(ROOT, "work", "stage6_stats.json"), "w") as f:
        json.dump(agg, f, indent=2)

    # concatenate parts (zstd frames concatenate losslessly)
    for k in range(N_BUCKETS):
        d = os.path.join(OUTDIR, "fold_%d" % k)
        pd = os.path.join(d, "parts")
        for split in ("train", "val", "test"):
            parts = sorted(glob.glob(os.path.join(pd, split + "-*.jsonl.zst")))
            outp = os.path.join(d, split + ".jsonl.zst")
            with open(outp, "wb") as out:
                for p in parts:
                    with open(p, "rb") as src:
                        while True:
                            chunk = src.read(1 << 24)
                            if not chunk:
                                break
                            out.write(chunk)
            for p in parts:
                os.remove(p)
        os.rmdir(pd)
        print("fold %d assembled" % k, flush=True)


if __name__ == "__main__":
    main()