23 lines
727 B
Python
23 lines
727 B
Python
|
"""
|
|||
|
|
BM25模型预下载脚本
|
|||
|
|
执行后将模型缓存到 ./models/fastembed_cache 目录，打包进Docker镜像
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 指定缓存目录
|
|||
|
|
cache_dir = "./models/fastembed_cache"
|
|||
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
print("正在下载BM25稀疏向量模型...")
|
|||
|
|
model = SparseTextEmbedding(
|
|||
|
|
model_name="Qdrant/bm25",
|
|||
|
|
cache_dir=cache_dir
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 触发一次推理，确保模型文件完整下载
|
|||
|
|
list(model.embed(["init trigger"]))
|
|||
|
|
print(f"✅ BM25模型已成功缓存到: {cache_dir}")
|
|||
|
|
print("请将该目录提交到项目仓库，打包进Docker镜像")
|