"""
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镜像")