83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
"""
|
|
下载 BGE embedding 和 reranker 模型到本地 models/ 目录
|
|
|
|
首次配置时运行：
|
|
python scripts/download_models.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
MODELS_DIR = os.path.join(PROJECT_ROOT, "models")
|
|
EMBEDDING_MODEL_PATH = os.path.join(MODELS_DIR, "bge-base-zh-v1.5")
|
|
RERANK_MODEL_PATH = os.path.join(MODELS_DIR, "bge-reranker-base")
|
|
|
|
os.makedirs(MODELS_DIR, exist_ok=True)
|
|
|
|
|
|
def download_embedding_model():
|
|
"""下载 BGE 中文向量模型"""
|
|
if os.path.exists(os.path.join(EMBEDDING_MODEL_PATH, "config.json")):
|
|
print(f"[跳过] Embedding 模型已存在: {EMBEDDING_MODEL_PATH}")
|
|
return True
|
|
|
|
print(f"[下载] BGE 中文向量模型 (bge-base-zh-v1.5)...")
|
|
print(f" 目标路径: {EMBEDDING_MODEL_PATH}")
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
model = SentenceTransformer("BAAI/bge-base-zh-v1.5")
|
|
model.save(EMBEDDING_MODEL_PATH)
|
|
print(f"[完成] Embedding 模型下载成功!")
|
|
return True
|
|
except Exception as e:
|
|
print(f"[失败] Embedding 模型下载失败: {e}")
|
|
return False
|
|
|
|
|
|
def download_reranker_model():
|
|
"""下载 BGE reranker 模型"""
|
|
if os.path.exists(os.path.join(RERANK_MODEL_PATH, "config.json")):
|
|
print(f"[跳过] Reranker 模型已存在: {RERANK_MODEL_PATH}")
|
|
return True
|
|
|
|
print(f"[下载] BGE Reranker 模型 (bge-reranker-base)...")
|
|
print(f" 目标路径: {RERANK_MODEL_PATH}")
|
|
try:
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
os.makedirs(RERANK_MODEL_PATH, exist_ok=True)
|
|
model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-base")
|
|
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-base")
|
|
model.save_pretrained(RERANK_MODEL_PATH)
|
|
tokenizer.save_pretrained(RERANK_MODEL_PATH)
|
|
print(f"[完成] Reranker 模型下载成功!")
|
|
return True
|
|
except Exception as e:
|
|
print(f"[失败] Reranker 模型下载失败: {e}")
|
|
print(f" (项目启动时会自动尝试下载，或可稍后手动重试)")
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 60)
|
|
print("RAG 项目模型下载脚本")
|
|
print("=" * 60)
|
|
|
|
ok1 = download_embedding_model()
|
|
print()
|
|
ok2 = download_reranker_model()
|
|
|
|
print()
|
|
print("=" * 60)
|
|
if ok1:
|
|
print("Embedding 模型: OK (项目可以启动)")
|
|
else:
|
|
print("Embedding 模型: FAILED (项目无法启动，请检查网络)")
|
|
|
|
if ok2:
|
|
print("Reranker 模型: OK")
|
|
else:
|
|
print("Reranker 模型: FAILED (不影响启动，首次使用 rerank 时自动下载)")
|
|
|
|
print("=" * 60)
|