#!/bin/bash
# Stage 2b: de-bias stage 2 with short-answer VQA data (continues from stage-2 LoRA + projector), then evaluate.
set -uo pipefail
cd "$HOME/atlas"
export PYTHONUNBUFFERED=1 HF_HOME="$HOME/.cache/huggingface" PATH="$HOME/.local/bin:$PATH"
echo "=== [1/4] build de-bias mix ==="
python3 prepare_mix.py || exit 11
export HF_HUB_OFFLINE=1
TRAIN_ENV="INSTRUCT_JSON=$HOME/data/llava_instruct/mix_debias.json INIT_FROM=$HOME/checkpoints/stage2/latest.pt LORA_LR=${LORA_LR:-1e-4} PROJ_LR=${PROJ_LR:-1e-5} HELDOUT=200"
echo "=== [2/4] smoke test (10 steps) ==="
rm -rf "$HOME/checkpoints/stage2b_smoke"
env $TRAIN_ENV MAX_STEPS=10 LOG_EVERY=2 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/stage2b_smoke" python3 train_stage2.py || { echo "SMOKE TEST FAILED"; exit 12; }
echo "=== [3/4] full stage-2b run ==="
env $TRAIN_ENV CKPT_DIR="$HOME/checkpoints/stage2b" python3 train_stage2.py || { echo "TRAINING FAILED"; exit 13; }
echo "=== [4/4] evaluation (same held-out set and POPE as stage 2) ==="
env -u INSTRUCT_JSON -u INIT_FROM -u HELDOUT CKPT_DIR="$HOME/checkpoints/stage2b" EVAL_DIR="$HOME/eval/stage2b" python3 eval_stage2.py || { echo "EVAL FAILED"; exit 14; }
cat "$HOME/eval/stage2b/report.md"
echo "=== run_stage2b.sh finished ==="