#!/bin/bash
set -euo pipefail

# ─── SmolLM3-3B LoRA Training Pipeline ───
# Stages: prepare → train (skip with env vars)
#
# Usage:
#   docker run --gpus all -v /path/to/output:/data smollora
#   docker run --gpus all -e SKIP_PREP=1 -v /path/to/processed:/data/processed smollora

MODEL="${MODEL:-HuggingFaceTB/SmolLM3-3B}"
DATA_DIR="${DATA_DIR:-/data}"
OUTPUT_DIR="${OUTPUT_DIR:-/output}"
EPOCHS="${EPOCHS:-3}"
BATCH_SIZE="${BATCH_SIZE:-4}"
LR="${LR:-2e-4}"
LORA_R="${LORA_R:-16}"
MAX_LENGTH="${MAX_LENGTH:-4096}"

echo "🎭 SmolLM3-3B LoRA Training Pipeline"
echo "   Model:    $MODEL"
echo "   Data:     $DATA_DIR"
echo "   Output:   $OUTPUT_DIR"
echo "   Epochs:   $EPOCHS"
echo "   Batch:    $BATCH_SIZE"
echo "   LR:       $LR"
echo "   LoRA r:   $LORA_R"
echo ""

# Stage 1: Data preparation
if [ "${SKIP_PREP:-0}" = "0" ]; then
    echo "━━━ Stage 1: Data Preparation ━━━"
    python /app/prepare_data.py \
        --output-dir "$DATA_DIR"
    echo "✅ Data prepared in $DATA_DIR"
else
    echo "⏭  Skipping data preparation (SKIP_PREP=1)"
fi

# Stage 2: Training
echo ""
echo "━━━ Stage 2: LoRA Training ━━━"
python /app/train_lora.py \
    --data-dir "$DATA_DIR" \
    --model "$MODEL" \
    --output-dir "$OUTPUT_DIR" \
    --epochs "$EPOCHS" \
    --batch-size "$BATCH_SIZE" \
    --lr "$LR" \
    --lora-r "$LORA_R" \
    --max-length "$MAX_LENGTH"

echo ""
echo "🎭 Training complete! Adapter saved to $OUTPUT_DIR/final"