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:
1 parent
a11f703c49
commit
49dd2be4cd
1 file changed
+9
-7
@@ -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,
|
def load_parquet_texts(path: str, kind: str, rng: random.Random,
|
||||||
cap: int) -> list[str]:
|
cap: int) -> list[str]:
|
||||||
import pyarrow.parquet as pq
|
import pyarrow.parquet as pq
|
||||||
out = []
|
d = pq.read_table(path).to_pydict()
|
||||||
pf = pq.ParquetFile(path)
|
|
||||||
for rg in pf.iter_row_groups():
|
|
||||||
d = rg.to_pydict()
|
|
||||||
n = len(next(iter(d.values())))
|
n = len(next(iter(d.values())))
|
||||||
idx = range(n) if n <= cap else rng.sample(range(n), cap)
|
idx = range(n) if n <= cap else rng.sample(range(n), cap)
|
||||||
|
out = []
|
||||||
for i in idx:
|
for i in idx:
|
||||||
if kind == "code":
|
if kind == "code":
|
||||||
|
if "content" in d:
|
||||||
|
frag = d["content"][i]
|
||||||
|
else:
|
||||||
frag = ("// %s\n### Instruction:\n%s\n### Response:\n%s"
|
frag = ("// %s\n### Instruction:\n%s\n### Response:\n%s"
|
||||||
% (d.get("file_name", [""])[i], d["prompt"][i],
|
% (d.get("file_name", [""])[i], d["prompt"][i],
|
||||||
d["code"][i] or d["response"][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]
|
frag = d["text"][i]
|
||||||
if frag and frag.strip():
|
if frag and frag.strip():
|
||||||
out.append(frag[:8000])
|
out.append(frag[:8000])
|
||||||
if len(out) >= cap:
|
|
||||||
break
|
|
||||||
rng.shuffle(out)
|
rng.shuffle(out)
|
||||||
return out[:cap]
|
return out[:cap]
|
||||||
|
|
||||||
@@ -211,11 +210,14 @@ def main() -> None:
|
|||||||
fr["kld"] = [text_of(p) for p in src["kld"]]
|
fr["kld"] = [text_of(p) for p in src["kld"]]
|
||||||
fr["cal"] = [text_of(p) for p in src["cal"]]
|
fr["cal"] = [text_of(p) for p in src["cal"]]
|
||||||
frags[cls] = fr
|
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),
|
for cls, spec in (("code", a.hf_code_parquet),
|
||||||
("prose", a.hf_prose_parquet)):
|
("prose", a.hf_prose_parquet)):
|
||||||
for path in [p for p in spec.split(",") if p]:
|
for path in [p for p in spec.split(",") if p]:
|
||||||
frags[cls]["cal"].extend(
|
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:
|
for cls in frags:
|
||||||
rng.shuffle(frags[cls]["cal"])
|
rng.shuffle(frags[cls]["cal"])
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user