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