#!/bin/bash
# Training + measurement data: diverse corpus rows (12/shard) -> teacher trajectories (train)
# + one pooled 128-row eval set. Selector written inline (make_traindata.py kept as reference).
set -euo pipefail
cd /root/ingot
export PATH=/root/ingot/.venv/bin:/usr/local/cuda/bin:$PATH
export INGOT_TEACHER_GPUS=0 INGOT_KING_GPUS=0 INGOT_CHALLENGER_GPUS=0
export INGOT_HIDDEN_STATES_DIR=/root/states-live
export HF_HUB_OFFLINE=1
PY=/root/ingot/.venv/bin/python

echo "== corpus rows -> /root/drawn.jsonl (12/shard across ~17 shards) =="
if [ ! -s /root/drawn.jsonl ]; then
$PY - <<'MK'
import json, os, urllib.request
from pathlib import Path
import pyarrow.parquet as pq
TOK = os.environ.get("HFTOKEN", "")
HDR = {"User-Agent": "mk", **({"Authorization": f"Bearer {TOK}"} if TOK else {})}

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=true",
    headers=HDR)
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:
    local = tmp / Path(shard).name
    if not local.exists():
        r = urllib.request.urlopen(urllib.request.Request(BASE + shard, headers=HDR))
        local.write_bytes(r.read())

    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()
    kept_shard = 0
    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 > 12000:
            continue
        sid = f"{row['dataset']}/{row['split']}/{row['entry_name']}/{row['language']}"
        out.append({"id": sid, "messages": msgs, "tools": tools})
        kept_shard += 1
        if kept_shard >= 12:
            break
    print(f"{shard}: kept {kept_shard}", flush=True)
    local.unlink(missing_ok=True)  # stream: /tmp can't hold all shards
    if len(out) >= 200:
        break

Path("/root/drawn.jsonl").write_text(
    "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in out[:200]))
print("wrote /root/drawn.jsonl:", min(len(out), 200), "rows from",
      len({r["id"].split("/")[0] for r in out}), "sources")
MK
fi

echo "== split: 48 train rows + one pooled 128-row eval set (disjoint) =="
$PY - <<'PY'
import json
rows = [json.loads(l) for l in open("/root/drawn.jsonl")][:176]
with open("/root/evalset.jsonl", "w") as f:
    for r in rows[48:176]:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
with open("/root/trainrows.jsonl", "w") as f:
    for r in rows[:48]:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
print("trainrows 48, evalset", min(128, max(0, len(rows) - 48)))
PY

echo "== teacher trajectories (vLLM, ~40-70 min) =="
[ -s /root/train_samples.json ] || $PY /root/gen_teacher.py --dataset /root/trainrows.jsonl \
  --samples 48 --max-prompt 1200 --workers 8 --out /root/train_samples.json
echo "GEN_DATA COMPLETE"