hako/sources/corpus_cache.py

"""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))