File size: 4,557 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 | """Build memmap shards from the raw parquet corpus.
Output layout (--out):
chars.bin boundary.bin dia.bin cap.bin aligned uint8 planes, one entry per letter
index.parquet per-record: offset, length, tier, clean, source, id
stats.json merged normalization Stats
Records are laid out pristine-files-first, then repaired, preserving parquet order,
so downstream passes (dedup keeps the first copy seen) prefer clean text.
"""
from __future__ import annotations
import argparse, json, sys
from collections import Counter
from concurrent.futures import ProcessPoolExecutor
from dataclasses import asdict
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 data.normalize import Stats, normalize_record
PLANES = ("chars", "boundary", "dia", "cap", "punct")
def process_file(args):
pf, tmpdir = args
pf, tmpdir = Path(pf), Path(tmpdir)
tmpdir.mkdir(parents=True, exist_ok=True)
t = pq.read_table(pf)
stats = Stats()
bufs = {p: [] for p in PLANES}
rows = {"offset": [], "length": [], "tier": [], "clean": [], "source": [], "id": []}
off = 0
texts = t.column("text").to_pylist()
tiers = t.column("tier").to_pylist()
cleans = t.column("clean").to_pylist()
sources = t.column("source").to_pylist()
ids = t.column("id").to_pylist()
for text, tier, clean, source, rid in zip(texts, tiers, cleans, sources, ids):
r = normalize_record(text, stats, with_punct=True)
if r is None:
continue
chars, boundary, dia, cap, punct = r
for p, a in zip(PLANES, r):
bufs[p].append(a)
rows["offset"].append(off)
rows["length"].append(len(chars))
rows["tier"].append(tier)
rows["clean"].append(clean)
rows["source"].append(source)
rows["id"].append(rid)
off += len(chars)
for p in PLANES:
np.concatenate(bufs[p]).tofile(tmpdir / f"{p}.bin")
pq.write_table(pa.table(rows), tmpdir / "index.parquet")
st = asdict(stats)
st["stripped"] = dict(stats.stripped)
st["archaic"] = dict(stats.archaic)
st["other_marks"] = dict(stats.other_marks)
(tmpdir / "stats.json").write_text(json.dumps(st))
return str(pf.name), off
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--raw", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--workers", type=int, default=12)
a = ap.parse_args()
raw, out = Path(a.raw), Path(a.out)
out.mkdir(parents=True, exist_ok=True)
files = sorted(raw.glob("pristine/*.parquet")) + sorted(raw.glob("repaired/*.parquet"))
tmp = out / "tmp"
jobs = [(str(f), str(tmp / f.stem)) for f in files]
with ProcessPoolExecutor(max_workers=a.workers) as ex:
for name, n in ex.map(process_file, jobs):
print(f" {name}: {n/1e6:.1f}M letters", flush=True)
# concatenate in canonical order
total = Stats()
tables, cum = [], 0
for f in files:
d = tmp / f.stem
t = pq.read_table(d / "index.parquet")
t = t.set_column(0, "offset", pa.array(t.column("offset").to_numpy() + cum))
tables.append(t)
cum += int(np.sum(t.column("length").to_numpy()))
s = json.loads((d / "stats.json").read_text())
st = Stats(**{k: s[k] for k in ("records_in", "records_kept",
"records_dropped_nongreek", "records_dropped_empty", "letters",
"words", "sentences", "mark_conflicts", "orphan_marks")})
st.stripped = Counter({int(k): v for k, v in s["stripped"].items()})
st.archaic = Counter(s["archaic"])
st.other_marks = Counter(s["other_marks"])
total.merge(st)
for p in PLANES:
with open(out / f"{p}.bin", "wb") as fo:
for f in files:
fo.write((tmp / f.stem / f"{p}.bin").read_bytes())
pq.write_table(pa.concat_tables(tables), out / "index.parquet")
st = asdict(total)
st["stripped"] = {str(k): v for k, v in total.stripped.items()}
st["archaic"] = dict(total.archaic)
st["other_marks"] = dict(total.other_marks)
(out / "stats.json").write_text(json.dumps(st, indent=2, ensure_ascii=False))
print(f"TOTAL letters: {cum/1e9:.3f}B records kept: {total.records_kept}/{total.records_in}")
print(f"words: {total.words/1e6:.1f}M sentences: {total.sentences/1e6:.1f}M")
if __name__ == "__main__":
main()
|