#!/bin/bash
# ============================================================
# VetNet DAPT 训练环境部署脚本 (hf-mirror)
# 服务器路径: /extra-storage/vetnet
# ============================================================
set -e

export HF_ENDPOINT=https://hf-mirror.com
BASE_DIR="/extra-storage/vetnet"
mkdir -p "$BASE_DIR"
cd "$BASE_DIR"

echo "============================================"
echo "VetNet DAPT 部署 — $(date)"
echo "目标: $BASE_DIR"
echo "============================================"

# ─── 1. 下载模型 ───
echo ""
echo "[1/3] 下载 Qwen2.5-3B-Instruct..."
MODEL_DIR="$BASE_DIR/Qwen2.5-3B-Instruct"
if [ -d "$MODEL_DIR" ] && [ -f "$MODEL_DIR/config.json" ]; then
    echo "  已存在，跳过"
else
    pip install huggingface_hub -q
    python3 -c "
from huggingface_hub import snapshot_download
snapshot_download('Qwen/Qwen2.5-3B-Instruct', local_dir='$MODEL_DIR', max_workers=1)
"
fi

# ─── 2. 下载训练包 ───
echo ""
echo "[2/3] 下载训练包..."
PKG="$BASE_DIR/vetnet_train_package.tar.gz"
if [ -f "$PKG" ]; then
    echo "  已存在，跳过"
else
    python3 -c "
from huggingface_hub import hf_hub_download
hf_hub_download('WWsCa/vetnet-train-package', 'vetnet_train_package.tar.gz', local_dir='$BASE_DIR')
"
fi

# ─── 3. 解压并安装 ───
echo ""
echo "[3/3] 解压 + 安装依赖..."
tar -xzf "$PKG" -C "$BASE_DIR" --overwrite
pip install -r "$BASE_DIR/requirements_gpu.txt" -q

echo ""
echo "============================================"
echo "部署完成！"
echo ""
echo "启动 DAPT 训练："
echo "  cd $BASE_DIR"
echo "  python train_full.py --dapt \\"
echo "      --model_path ./Qwen2.5-3B-Instruct \\"
echo "      --train_data ./train_data.jsonl \\"
echo "      --val_data ./train_data_val.jsonl \\"
echo "      --output_dir ./vet-qwen3b-v3-dapt \\"
echo "      --epochs 5 \\"
echo "      --batch_size 4 \\"
echo "      --gradient_accumulation 8 \\"
echo "      --max_length 4096 \\"
echo "      --learning_rate 5e-5"
echo "============================================"