#!/bin/bash
# Qwen3-Omni Thinker 训练快速启动脚本

set -e  # 遇到错误立即退出

echo "=========================================="
echo "Qwen3-Omni Thinker 训练环境检查"
echo "=========================================="

# 1. 检查Python环境
echo "检查Python版本..."
python --version

# 2. 检查必要的包
echo ""
echo "检查必要的Python包..."
REQUIRED_PACKAGES=("torch" "transformers" "datasets" "deepspeed" "accelerate")

for package in "${REQUIRED_PACKAGES[@]}"; do
    if python -c "import $package" 2>/dev/null; then
        echo "✓ $package 已安装"
    else
        echo "✗ $package 未安装"
        echo "  安装命令: pip install $package"
    fi
done

# 3. 检查CUDA
echo ""
echo "检查CUDA环境..."
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --query-gpu=name,memory.total --format=csv
else
    echo "警告: nvidia-smi 未找到"
fi

# 4. 检查必要的文件
echo ""
echo "检查配置文件..."
FILES=("train.py" "configs/ds_config_zero2.json" "train_qwen3_omni.sbatch")

for file in "${FILES[@]}"; do
    if [ -f "$file" ]; then
        echo "✓ $file 存在"
    else
        echo "✗ $file 不存在"
    fi
done

# 5. 创建必要的目录
echo ""
echo "创建必要的目录..."
mkdir -p logs
mkdir -p output
mkdir -p data
echo "✓ 目录创建完成"

# 6. 检查模型和数据路径配置
echo ""
echo "=========================================="
echo "下一步操作："
echo "=========================================="
echo "1. 下载或准备Qwen3-Omni模型"
echo "   例如: git lfs clone https://huggingface.co/Qwen/Qwen3-Omni-base"
echo ""
echo "2. 准备训练数据"
echo "   方式1: 创建示例数据测试"
echo "   $ python prepare_data.py --create-sample --output data/train.jsonl"
echo ""
echo "   方式2: 转换您的数据"
echo "   $ python prepare_data.py --input your_data.jsonl --output data/train.jsonl"
echo ""
echo "3. 修改SLURM脚本中的路径"
echo "   编辑 train_qwen3_omni.sbatch 或 train_qwen3_omni_single.sbatch"
echo "   - MODEL_PATH: 模型路径"
echo "   - DATA_PATH: 数据路径"
echo "   - SBATCH参数: 根据集群配置修改"
echo ""
echo "4. 提交训练任务"
echo "   多节点: sbatch train_qwen3_omni.sbatch"
echo "   单节点: sbatch train_qwen3_omni_single.sbatch"
echo ""
echo "5. 监控训练"
echo "   查看队列: squeue -u \$USER"
echo "   查看日志: tail -f logs/train_<JOB_ID>.log"
echo "=========================================="