34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
# 本文件用于下载并获取embedding向量化模型
|
|
|
|
from pathlib import Path
|
|
from langchain_huggingface import HuggingFaceEmbeddings
|
|
|
|
# 获取embeddings模型 - 首次调用时自动下载
|
|
def get_embeddings(model_name="BAAI/bge-small-zh-v1.5",device="cpu",**kwargs):
|
|
# 支持更换其他向量化模型
|
|
local_dir = Path("models")/model_name.replace("/", "_")
|
|
if not local_dir.exists():
|
|
print(f'⚠️ 首次使用嵌入模型，正在下载到{local_dir.absolute()}')
|
|
print("💡 提示：需要联网(必需梯子)，完成后可离线使用")
|
|
from huggingface_hub import snapshot_download
|
|
# 模型下载工具
|
|
snapshot_download(
|
|
repo_id = model_name,
|
|
local_dir = local_dir,
|
|
)
|
|
print('✅ 下载完成！')
|
|
|
|
# 构造参数字典
|
|
model_kwargs = {
|
|
"device":device,
|
|
"local_files_only": True, # 仅使用本地文件
|
|
}
|
|
|
|
# 此处实例化时，把kwargs传入
|
|
_EMBEDDINGS = HuggingFaceEmbeddings(
|
|
model_name = str(local_dir), # 使用本地已下载的模型
|
|
model_kwargs = model_kwargs,
|
|
**kwargs # 允许传入参数
|
|
)
|
|
return _EMBEDDINGS
|