rag_jrxml/down_embedding_model.py

"""
down_embedding_model.py
下载嵌入模型（模型名称通过 .env / config.py 配置）
"""

import os
import sys
from pathlib import Path
from config import EMBEDDING_MODEL_NAME, EMBEDDING_MODEL_PATH, HF_ENDPOINT

def download_model():
    """下载嵌入模型"""
    model_dir = EMBEDDING_MODEL_PATH

    print("=" * 60)
    print(f"{EMBEDDING_MODEL_NAME} 模型下载")
    print("=" * 60)
    print(f"模型名称: {EMBEDDING_MODEL_NAME}")
    print(f"模型目录: {model_dir}")
    print()

    os.environ['HF_ENDPOINT'] = HF_ENDPOINT
    print(f"使用 HuggingFace 镜像: {HF_ENDPOINT}")
    print()

    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("❌ 未安装 huggingface_hub，正在安装...")
        import subprocess
        subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub"])
        from huggingface_hub import snapshot_download

    # 创建模型目录
    os.makedirs(model_dir, exist_ok=True)

    print(f"开始下载 {EMBEDDING_MODEL_NAME} 模型...")
    print(f"请耐心等待...")
    print()

    try:
        snapshot_download(
            repo_id=EMBEDDING_MODEL_NAME,
            local_dir=str(model_dir),
            local_dir_use_symlinks=False,
            resume_download=True
        )
        print()
        print("=" * 60)
        print("✅ 模型下载完成！")
        print("=" * 60)
        print(f"模型路径: {model_dir}")
        return True
    except Exception as e:
        print()
        print("=" * 60)
        print(f"❌ 下载失败: {e}")
        print("=" * 60)
        return False

if __name__ == "__main__":
    download_model()