push artifacts/scripts/make_traindata.py

artifacts/scripts/make_traindata.py

@@ -0,0 +1,59 @@
+"""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")