#!/bin/bash
# 训练启动脚本 - 使用 LLaMA-Factory 进行 QLoRA 训练
# 适用于 Qwen3.5-0.8B-HOS 训练流程

set -e

echo "============================================"
echo "  启动 QLoRA 训练"
echo "  模型: Qwen3.5-0.8B-HOS"
echo "============================================"

# 检查配置文件
CONFIG_FILE="configs/train.yaml"
if [ ! -f "$CONFIG_FILE" ]; then
    echo "错误: 配置文件不存在: $CONFIG_FILE"
    exit 1
fi

# 检查数据集
DATASET_DIR="datasets"
if [ ! -d "$DATASET_DIR" ]; then
    echo "错误: 数据集目录不存在: $DATASET_DIR"
    exit 1
fi

# 创建输出目录
OUTPUT_DIR="outputs/qwen35-0.8b-cybersec-qlora"
mkdir -p "$OUTPUT_DIR"

# 显示训练配置
echo "训练配置:"
echo "  配置文件: $CONFIG_FILE"
echo "  数据集目录: $DATASET_DIR"
echo "  输出目录: $OUTPUT_DIR"
echo ""

# 显示系统信息
echo "系统信息:"
python -c "import torch; print(f'  PyTorch: {torch.__version__}')"
python -c "import torch; print(f'  CUDA Available: {torch.cuda.is_available()}')"
if python -c "import torch; exit(0 if torch.cuda.is_available() else 1)" 2>/dev/null; then
    python -c "import torch; print(f'  GPU: {torch.cuda.get_device_name(0)}')"
    python -c "import torch; print(f'  GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB')"
fi
echo ""

# 开始训练
echo "开始训练..."
echo "============================================"

llamafactory-cli train "$CONFIG_FILE"

echo "============================================"
echo "  训练完成!"
echo "  模型输出: $OUTPUT_DIR"
echo "============================================"