#!/bin/bash
# 环境安装脚本 - 安装训练所需的所有依赖
# 适用于 Qwen3.5-0.8B-HOS 训练流程

set -e

echo "============================================"
echo "  安装训练环境依赖"
echo "============================================"

# 检查 Python 版本
echo "[1/6] 检查 Python 版本..."
python --version

# 升级 pip
echo "[2/6] 升级 pip..."
python -m pip install --upgrade pip

# 安装 PyTorch (CUDA 12.1)
echo "[3/6] 安装 PyTorch..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 安装 LLaMA-Factory
echo "[4/6] 安装 LLaMA-Factory..."
pip install git+https://github.com/hiyouga/LLaMA-Factory.git

# 安装 Unsloth (加速训练)
echo "[5/6] 安装 Unsloth..."
pip install unsloth

# 安装 vLLM (推理加速)
echo "[6/6] 安装 vLLM..."
pip install vllm

# 安装其他依赖
echo "[7/7] 安装其他依赖..."
pip install -r requirements.txt

# 验证安装
echo "============================================"
echo "  验证安装..."
echo "============================================"
python -c "import torch; print(f'PyTorch: {torch.__version__}')"
python -c "import torch; print(f'CUDA Available: {torch.cuda.is_available()}')"
python -c "import torch; print(f'CUDA Version: {torch.version.cuda}')"
python -c "import transformers; print(f'Transformers: {transformers.__version__}')"

echo "============================================"
echo "  环境安装完成!"
echo "============================================"