#!/bin/bash
# stage2: depth pruning of f28 to 16/12/8 layers, retrain on the full set, evaluate
set -e
cd $HOME/mt92
export TOKENIZERS_PARALLELISM=false
for K in 16 12 8; do
  EP=4; [ $K -le 12 ] && EP=6; [ $K -le 8 ] && EP=8
  echo "##### keep $K layers (epochs $EP)"
  [ -d p$K ] || ./venv/bin/python prune.py f28 p$K $K 2>&1 | grep -E 'influence|dropping|saved'
  [ -d r$K ] || OUT=r$K BASE=p$K EPOCHS=$EP LR=8e-5 TRAIN=ex/full.jsonl ./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 r$K ex/$s.jsonl; done
done
echo STAGE2_DONE