#!/bin/bash
# VLM预训练启动脚本

# 设置环境变量
export CUDA_VISIBLE_DEVICES=0
export PYTHONPATH="PROJECT_ROOT:$PYTHONPATH"

# 训练目录
TRAIN_DIR="PROJECT_ROOT/training/pretrain"
mkdir -p $TRAIN_DIR
cd $TRAIN_DIR

echo "======================================"
echo "VLM预训练"
echo "======================================"
echo ""

# 检查数据
DATA_FILE="PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data.pkl"
if [ ! -f "$DATA_FILE" ]; then
    echo "❌ 数据文件不存在: $DATA_FILE"
    echo "请先运行: python prepare_pretrain_data.py"
    exit 1
fi
echo "✓ 数据文件: $DATA_FILE"

# 创建输出目录
OUTPUT_DIR="PROJECT_ROOT/checkpoints/pretrain"
mkdir -p $OUTPUT_DIR
echo "✓ 输出目录: $OUTPUT_DIR"
echo ""

# 选择模型
MODEL=$1
if [ -z "$MODEL" ]; then
    echo "用法: bash run_pretrain.sh [qwen2.5-vl-3b|qwen2.5-vl-7b]"
    echo ""
    echo "示例:"
    echo "  bash run_pretrain.sh qwen2.5-vl-3b"
    echo "  bash run_pretrain.sh qwen2.5-vl-7b"
    exit 1
fi

echo "======================================"
echo "开始训练: $MODEL"
echo "======================================"
echo ""

# 训练
python train_pretrain.py \
    --model $MODEL \
    --epochs 5 \
    --batch_size 1 \
    --lr 2e-5

echo ""
echo "======================================"
echo "训练完成！"
echo "======================================"
echo "Checkpoints: $OUTPUT_DIR/$MODEL"