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:
@@ -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"])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user