progress / artifacts /scripts /make_traindata.py
eyad-silx's picture
push artifacts/scripts/make_traindata.py
3d5a84d verified
Raw History Blame
2.23 kB
"""Quick training-data builder: a few corpus shards at the eval-156 commit, rows with short
prompts, dumped as drawn.jsonl (the dataset.load format). No ranking needed: the QAT trains on
any diverse rows; the margin is verified on the official harness later."""
import json
import sys
import urllib.request
from pathlib import Path
sys.path.insert(0, "/root/ingot/src")
import pyarrow.parquet as pq
COMMIT = "70d2121e5dcbfc9f2c2e79c890563d96429eeb98"
REPO = "dendriteholdings/ingot_corpus"
BASE = f"https://huggingface.co/datasets/{REPO}/resolve/{COMMIT}/"
req = urllib.request.Request(
f"https://huggingface.co/api/datasets/{REPO}/tree/{COMMIT}?recursive=1",
headers={"User-Agent": "mk"},
)
tree = json.load(urllib.request.urlopen(req))
shards = sorted(
t["path"] for t in tree if t["path"].startswith("data/chungus-") and t["path"].endswith(".parquet")
)
print("shards:", len(shards))
out = []
tmp = Path("/tmp/shards")
tmp.mkdir(exist_ok=True)
for shard in shards[::17]: # ~15 shards, spread over the corpus
url = BASE + shard
local = tmp / Path(shard).name
urllib.request.urlretrieve(url, local)
parq = pq.ParquetFile(local)
cols = ["dataset", "split", "entry_name", "domain", "language", "user_prompt"]
if "conversation" in parq.schema_arrow.names:
cols.append("conversation")
df = parq.read(columns=cols).to_pylist()
for row in df:
if row.get("conversation"):
conv = json.loads(row["conversation"])
msgs, tools = conv["messages"], conv.get("tools")
else:
msgs = [{"role": "user", "content": row["user_prompt"]}]
tools = None
size = sum(len(m["content"].encode()) for m in msgs)
if size > 3500: # short-ish prompts only
continue
sid = f"{row['dataset']}/{row['split']}/{row['entry_name']}/{row['language']}"
out.append({"id": sid, "messages": msgs, "tools": tools})
if len(out) >= 160:
break
print(f"{shard}: {len(out)} kept", flush=True)
if len(out) >= 160:
break
Path("/root/drawn.jsonl").write_text(
"".join(json.dumps(r, ensure_ascii=False) + "\n" for r in out)
)
print("wrote /root/drawn.jsonl:", len(out), "rows")