#!/bin/bash
# stage5: slate models from the train-only base m_lr8
#   c28: 28 layers, 60% of val/test     -> hotkey 1/1   (uid 169)
#   b16: 16 layers, 35% of val/test     -> hotkey hata/4 (uid 52)
#   d8 :  8 layers, 70% of val/test     -> hotkey hata/5 (uid 53)
set -e
cd $HOME/mt92
export TOKENIZERS_PARALLELISM=false
set -a; . ./.env; set +a
curl -s -L https://huggingface.co/dasLOL/mt-extract-3g/resolve/main/manifest.json -o manifest_template.json
[ -d p16m ] || ./venv/bin/python prune.py m_lr8 p16m 16 2>&1 | grep -E 'dropping|saved'
[ -d p8m ]  || ./venv/bin/python prune.py m_lr8 p8m 8   2>&1 | grep -E 'dropping|saved'
train() { # name base epochs set
  echo "##### $1: base $2, $3 epochs, $4"
  [ -d $1 ] || OUT=$1 BASE=$2 EPOCHS=$3 LR=8e-5 TRAIN=$4 ./venv/bin/python sft.py 2>&1 | tr '\r' '\n' | grep -E "train_loss|saved|Error|error"
  for s in val_all test_all gen_holdout; do ./venv/bin/python evalset.py $1 ex/$s.jsonl; done
  ./venv/bin/python backup_now.py $1 > backup_$1.log 2>&1 && echo "$1 backed up"
  bash ./build_final.sh $1 32 > build_$1.log 2>&1 && grep -E "MICRO|artifact|TTFT" build_$1.log
  cp manifest_template.json onnx_$1/final/manifest.json
  ./venv/bin/python calib.py onnx_$1/final 60 > calib_$1.log 2>&1; grep -v "^probe" calib_$1.log | tr -d "\n"; echo
  ./venv/bin/python up_art.py onnx_$1/final onnx_$1 > up_$1.log 2>&1 && echo "staged onnx_$1"
}
train c28 m_lr8 4 ex/full60.jsonl
train b16 p16m 4 ex/full35.jsonl
train d8 p8m 8 ex/full70.jsonl
echo STAGE5_DONE