92 lines
3.6 KiB
Bash
92 lines
3.6 KiB
Bash
|
#!/usr/bin/env bash
|
|||
|
|
# ══════════════════════════════════════════════════
|
|||
|
|
# download_models.sh
|
|||
|
|
# 预下载 AI 模型到 ./models 目录（加速容器启动）
|
|||
|
|
# 支持 HuggingFace 镜像加速（国内网络）
|
|||
|
|
# 用法：bash scripts/download_models.sh
|
|||
|
|
# ══════════════════════════════════════════════════
|
|||
|
|
set -euo pipefail
|
|||
|
|
|
|||
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|||
|
|
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
|
|||
|
|
cd "$PROJECT_DIR"
|
|||
|
|
|
|||
|
|
BLUE='\033[0;34m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
|
|||
|
|
info() { echo -e "${BLUE}[INFO]${NC} $*"; }
|
|||
|
|
ok() { echo -e "${GREEN}[OK]${NC} $*"; }
|
|||
|
|
warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
|
|||
|
|
|
|||
|
|
MODELS_DIR="$PROJECT_DIR/models"
|
|||
|
|
mkdir -p "$MODELS_DIR"
|
|||
|
|
|
|||
|
|
# 设置镜像加速
|
|||
|
|
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
|
|||
|
|
export HF_HOME="$MODELS_DIR"
|
|||
|
|
info "HuggingFace 镜像：$HF_ENDPOINT"
|
|||
|
|
info "模型保存路径：$MODELS_DIR"
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
# ── 方法1：通过 huggingface_hub 下载 ────────────
|
|||
|
|
download_hf() {
|
|||
|
|
local repo=$1; local local_name=$2
|
|||
|
|
info "下载 $repo..."
|
|||
|
|
if python3 -c "
|
|||
|
|
import os
|
|||
|
|
os.environ['HF_ENDPOINT'] = '${HF_ENDPOINT}'
|
|||
|
|
os.environ['HF_HOME'] = '${MODELS_DIR}'
|
|||
|
|
from huggingface_hub import snapshot_download
|
|||
|
|
snapshot_download(repo_id='$repo', cache_dir='${MODELS_DIR}')
|
|||
|
|
print('下载完成')
|
|||
|
|
" 2>&1; then
|
|||
|
|
ok "$repo 下载成功"
|
|||
|
|
else
|
|||
|
|
warn "$repo HuggingFace 下载失败，尝试 ModelScope..."
|
|||
|
|
download_modelscope "$repo" "$local_name"
|
|||
|
|
fi
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ── 方法2：通过 ModelScope 下载（备用）──────────
|
|||
|
|
download_modelscope() {
|
|||
|
|
local hf_name=$1
|
|||
|
|
local ms_name=${2:-$1}
|
|||
|
|
python3 -c "
|
|||
|
|
try:
|
|||
|
|
from modelscope import snapshot_download
|
|||
|
|
snapshot_download(model_id='$ms_name', cache_dir='${MODELS_DIR}/modelscope')
|
|||
|
|
print('ModelScope 下载完成')
|
|||
|
|
except ImportError:
|
|||
|
|
print('ModelScope 未安装，跳过')
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f'ModelScope 下载失败: {e}')
|
|||
|
|
" 2>&1 || warn "ModelScope 下载也失败，模型将在容器启动时自动下载"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ── 检查 Python 环境 ────────────────────────────
|
|||
|
|
if ! python3 -c "import huggingface_hub" 2>/dev/null; then
|
|||
|
|
warn "未安装 huggingface_hub，尝试安装..."
|
|||
|
|
pip3 install -q huggingface_hub modelscope 2>/dev/null || \
|
|||
|
|
warn "安装失败，模型将在容器首次启动时下载"
|
|||
|
|
fi
|
|||
|
|
|
|||
|
|
# ── 下载模型列表 ────────────────────────────────
|
|||
|
|
info "=== 下载 BGE-M3 嵌入模型（约 2.5GB）==="
|
|||
|
|
download_hf "BAAI/bge-m3" "BAAI/bge-m3"
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
info "=== 下载 BGE-Reranker 精排模型（约 1.1GB）==="
|
|||
|
|
download_hf "BAAI/bge-reranker-v2-m3" "BAAI/bge-reranker-v2-m3"
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
# MinerU 模型通过容器内脚本下载（依赖 magic-pdf 配置）
|
|||
|
|
info "=== MinerU 模型说明 ==="
|
|||
|
|
warn "MinerU 模型（约 2GB）将在 mcp-server 容器首次启动时自动下载"
|
|||
|
|
warn "如需预下载，请在 mcp-server 容器内运行：mineru-models-download"
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
echo -e "${GREEN}══════════════════════════════════════════${NC}"
|
|||
|
|
echo -e "${GREEN} 模型下载完成！${NC}"
|
|||
|
|
echo -e "${GREEN}══════════════════════════════════════════${NC}"
|
|||
|
|
echo ""
|
|||
|
|
echo "已下载到：$MODELS_DIR"
|
|||
|
|
du -sh "$MODELS_DIR" 2>/dev/null || true
|