#!/bin/bash
# usage: run.sh <tag> [train args...]: data prepared once, then text and joint phases on one schedule of STEPS steps (default 17000)
set -e
source env.sh
TAG=$1; shift; STEPS=${STEPS:-17000}
P=${MODEL:-pruned}
MIX="smoltalk:0.47,everyday:0.01,apigen:0.10,hermes_glaive:0.02,hermes_fc:0.005,hermes_fc_single:0.005,nemotron:0.16,synth:0.225,ci_shapes:0.015"
[ -d tokdata ] || { python3 render.py $P; python3 synth_tools.py $P 300000 data/rendered/synth; python3 tokenize_data.py $P tokdata; }
COMMON="--model $P --bs 32 --seq 2048 --mix $MIX --steps $STEPS --warmup 200 --lr-muon 2e-3 --lr-adam 1e-3 --eval-every 500 --gen-every 1000 --save-every 2000"
python3 train.py $COMMON --out ckpt/$TAG-text --name $TAG-text --stop 6000 "$@" > train_$TAG-text.log 2>&1
python3 train.py $COMMON --out ckpt/$TAG-joint --name $TAG-joint --resume ckpt/$TAG-text/resume.pt --ocr-frac 0.35 --ocr-workers 12 "$@" > train_$TAG-joint.log 2>&1