File size: 1,821 Bytes
daa13bc | 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 | """Corpus cache extractor: materialize a small JSONL sample of the giant
single-row-group parquet (BrunoN-Dev/corpus-ptbr-v1) in an ISOLATED process.
Why: the shard has ONE row group of 500k rows (~1.2 GiB decompressed text).
Any in-process read (datasets streaming, iter_batches) inflates the resident
set of the trainer beyond the RAM cap. The P-SRC/P-MEM fix: run this module
as a STANDALONE python process (no torch, no onnx loaded), read the row
group once, write N docs to a JSONL cache, exit. The trainer then streams
from the ~10 MiB cache file.
Usage: python3 -m hako.sources.corpus_cache <out_jsonl> <n_docs>
"""
from __future__ import annotations
import json
import sys
def main(out_path: str, n_docs: int = 6000) -> int:
import json
import resource
import polars as pl
# Polars' Rust parquet reader decodes page-by-page with slice pushdown,
# so head(N) never materializes the 1-row-group column (~3 GiB in
# pyarrow). Peak RSS measured: ~2.4 GiB for 6000 docs.
uri = ("https://huggingface.co/datasets/BrunoN-Dev/corpus-ptbr-v1/"
"resolve/main/data/train-00000-of-00018.parquet")
df = pl.scan_parquet(uri).select("text").head(n_docs).collect(
engine="streaming")
texts = []
for t in df["text"].to_list():
t = str(t)[:2000].strip()
if len(t) > 40:
texts.append(t)
del df
with open(out_path, "w", encoding="utf-8") as fh:
for t in texts:
fh.write(json.dumps({"text": t}, ensure_ascii=True) + "\n")
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
print(f"cache written: {len(texts)} docs -> {out_path} "
f"(peak RSS {rss:.0f} MiB)")
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 6000))
|