From 49dd2be4cd94e91910e4a20cdea985103b8753cd Mon Sep 17 00:00:00 2001
From: Frank <frank@dankware>
Date: Tue, 6 Oct 2026 21:16:16 -0700
Subject: [PATCH] corpus: handle codeparrot content column + per-class caps

- load_parquet_texts: read codeparrot-clean 'content' column (schema-robust)
- per-class HF caps (prose 60k for short ~66-tok rows, code 40k for long files)
- pyarrow 18: ParquetFile.iter_row_groups removed -> read_table().to_pydict()

Produces 10,014,932 calib tokens (4.5M code/3.0M prose/2.5M tools) and
1,800,816 KLD tokens (197 prompts, all >=8192, score_from=n-256) from real
HF sources (codeparrot-clean + wikitext-103).
---
 scripts/make-calib-corpus.py | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/scripts/make-calib-corpus.py b/scripts/make-calib-corpus.py
index 86203a3..e6cf6eb 100755
--- a/scripts/make-calib-corpus.py
+++ b/scripts/make-calib-corpus.py
@@ -122,23 +122,22 @@ def tool_call(rng: random.Random, cmds: list[str]) -> str:
 def load_parquet_texts(path: str, kind: str, rng: random.Random,
                        cap: int) -> list[str]:
     import pyarrow.parquet as pq
+    d = pq.read_table(path).to_pydict()
+    n = len(next(iter(d.values())))
+    idx = range(n) if n <= cap else rng.sample(range(n), cap)
     out = []
-    pf = pq.ParquetFile(path)
-    for rg in pf.iter_row_groups():
-        d = rg.to_pydict()
-        n = len(next(iter(d.values())))
-        idx = range(n) if n <= cap else rng.sample(range(n), cap)
-        for i in idx:
-            if kind == "code":
+    for i in idx:
+        if kind == "code":
+            if "content" in d:
+                frag = d["content"][i]
+            else:
                 frag = ("// %s\n### Instruction:\n%s\n### Response:\n%s"
                         % (d.get("file_name", [""])[i], d["prompt"][i],
                            d["code"][i] or d["response"][i]))
-            else:
-                frag = d["text"][i]
-            if frag and frag.strip():
-                out.append(frag[:8000])
-        if len(out) >= cap:
-            break
+        else:
+            frag = d["text"][i]
+        if frag and frag.strip():
+            out.append(frag[:8000])
     rng.shuffle(out)
     return out[:cap]
 
@@ -211,11 +210,14 @@ def main() -> None:
             fr["kld"] = [text_of(p) for p in src["kld"]]
             fr["cal"] = [text_of(p) for p in src["cal"]]
         frags[cls] = fr
+    # per-class fragment caps: prose rows are short (~66 tok) so need many;
+    # codeparrot files are long (~2.6k tok) so a small cap suffices
+    hfcap = {"code": 40000, "prose": 60000, "tools": 0}
     for cls, spec in (("code", a.hf_code_parquet),
                       ("prose", a.hf_prose_parquet)):
         for path in [p for p in spec.split(",") if p]:
             frags[cls]["cal"].extend(
-                load_parquet_texts(path, cls, rng, cap=40000))
+                load_parquet_texts(path, cls, rng, cap=hfcap[cls]))
     for cls in frags:
         rng.shuffle(frags[cls]["cal"])