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).
This commit is contained in:
Frank
2026-10-06 21:16:16 -07:00
parent a11f703c49
commit 49dd2be4cd
+9 -7
View File
@@ -122,14 +122,15 @@ 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
out = []
pf = pq.ParquetFile(path)
for rg in pf.iter_row_groups():
d = rg.to_pydict()
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 = []
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]))
@@ -137,8 +138,6 @@ def load_parquet_texts(path: str, kind: str, rng: random.Random,
frag = d["text"][i]
if frag and frag.strip():
out.append(frag[:8000])
if len(out) >= cap:
break
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"])