34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
"""
|
||
BM25 稀疏嵌入器
|
||
基于 FastEmbed 的 Qdrant/bm25 模型，完全离线运行
|
||
"""
|
||
from typing import List
|
||
from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
|
||
from app.config import FASTEMBED_CACHE_PATH
|
||
|
||
class BM25SparseEmbedder:
|
||
"""BM25 稀疏嵌入包装器，与现有嵌入器风格统一"""
|
||
|
||
def __init__(self):
|
||
self.model = SparseTextEmbedding(
|
||
model_name="Qdrant/bm25",
|
||
cache_dir=FASTEMBED_CACHE_PATH,
|
||
local_files_only=True, # 强制离线，永不联网
|
||
)
|
||
|
||
def embed_documents(self, texts: List[str]) -> List[dict]:
|
||
"""返回稀疏向量列表，每个为 Qdrant 兼容的 dict（indices+values）"""
|
||
return [vec.as_object() for vec in self.model.embed(texts)]
|
||
|
||
def embed_query(self, text: str) -> dict:
|
||
"""返回单个稀疏向量"""
|
||
return list(self.model.embed([text]))[0].as_object()
|
||
|
||
# 全局单例
|
||
_sparse_embedder_instance = None
|
||
|
||
def get_sparse_embedder() -> BM25SparseEmbedder:
|
||
global _sparse_embedder_instance
|
||
if _sparse_embedder_instance is None:
|
||
_sparse_embedder_instance = BM25SparseEmbedder()
|
||
return _sparse_embedder_instance |