#!/bin/bash
# ─────────────────────────────────────────────────────────
# train.sh — Launch CXR VLM training (both stages)
# Usage: bash scripts/train.sh
# ─────────────────────────────────────────────────────────

# Exit on error
set -e

# Set project root
PROJ_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$PROJ_DIR"

# ── Config ────────────────────────────────────────────────
MODEL_CONFIG="configs/model_config.yaml"
TRAIN_CONFIG="configs/train_config.yaml"

# ── Optional: resume from checkpoint ──────────────────────
RESUME_FROM=""   # e.g. "checkpoints/stage2_instruct/checkpoint-500"

# ── WandB ─────────────────────────────────────────────────
export WANDB_PROJECT="cxr-vlm-thesis"

# ── Run both stages ───────────────────────────────────────
echo "=========================================="
echo " CXR VLM Training"
echo " Model config: $MODEL_CONFIG"
echo " Train config: $TRAIN_CONFIG"
echo "=========================================="

if [ -n "$RESUME_FROM" ]; then
    python -m training.train \
        --model_config "$MODEL_CONFIG" \
        --train_config "$TRAIN_CONFIG" \
        --resume_from  "$RESUME_FROM"
else
    python -m training.train \
        --model_config "$MODEL_CONFIG" \
        --train_config "$TRAIN_CONFIG"
fi

echo "Training complete!"