#!/usr/bin/env bash
# Launch SFT from pretrain_v2 Stage-B checkpoint.
# Usage:  bash training/SFT/train_sft.sh [--debug]
set -euo pipefail

ROOT=PROJECT_ROOT

PRETRAINED_LORA="$ROOT/checkpoints/pretrain_v2/stage_b/best_model"
MODEL_PATH="$ROOT/models/Qwen2.5-VL-3B-Instruct"
NEXAR_ROOT="$ROOT/NEXAR_COLLISION/dataset"
DADA_ROOT="$ROOT/DADA-2000"
OUTPUT_DIR="$ROOT/checkpoints/SFT"
EXPERIMENT="sft_from_pretrain_v2"

DEBUG_FLAGS=""
if [[ "${1:-}" == "--debug" ]]; then
    DEBUG_FLAGS="--debug --debug_samples 50"
    echo "=== DEBUG MODE ==="
fi

echo "Starting SFT training..."
echo "  Pretrained LoRA : $PRETRAINED_LORA"
echo "  Output          : $OUTPUT_DIR/$EXPERIMENT"

python -m training.SFT.trainer \
    --nexar_root "$NEXAR_ROOT" \
    --dada_root  "$DADA_ROOT"  \
    --model_name "$MODEL_PATH" \
    --pretrained_lora "$PRETRAINED_LORA" \
    --output_dir "$OUTPUT_DIR" \
    --experiment_name "$EXPERIMENT" \
    --num_epochs 10 \
    --batch_size 1 \
    --gradient_accumulation_steps 8 \
    --learning_rate 1e-4 \
    --tta_head_lr  1e-3 \
    --vlm_lr_multiplier 0.1 \
    --weight_decay 0.01 \
    --max_grad_norm 1.0 \
    --mse_weight 1.0 \
    --nll_weight 0.5 \
    --use_curriculum \
    --no_auto_resume \
    --use_wandb \
    $DEBUG_FLAGS