31 lines
859 B
Python
31 lines
859 B
Python
import logging
|
|
import os
|
|
from functools import lru_cache
|
|
|
|
from .schemas import SparseVector
|
|
|
|
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
|
FASTEMBED_CACHE_PATH = "/models/fastembed"
|
|
|
|
logger = logging.getLogger("index-service")
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_sparse_model():
|
|
from fastembed import SparseTextEmbedding
|
|
|
|
logger.info("Loading sparse model %s from cache %s", SPARSE_MODEL_NAME, FASTEMBED_CACHE_PATH)
|
|
return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
|
|
|
|
|
|
def embed_sparse_texts(texts: list[str]) -> list[SparseVector]:
|
|
model = get_sparse_model()
|
|
result: list[SparseVector] = []
|
|
for item in model.embed(texts):
|
|
result.append(
|
|
SparseVector(
|
|
indices=[int(i) for i in item.indices.tolist()],
|
|
values=[float(v) for v in item.values.tolist()],
|
|
)
|
|
)
|
|
return result
|