HAKO upload: hako/sources/corpus_cache.py

hako/sources/corpus_cache.py

@@ -0,0 +1,48 @@
+"""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))