109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
import os
|
|
import re
|
|
from functools import lru_cache
|
|
|
|
import httpx
|
|
from fastembed import SparseTextEmbedding
|
|
|
|
from config import (
|
|
EMBEDDINGS_DENSE_MODEL,
|
|
EMBEDDINGS_DENSE_URL,
|
|
SPARSE_MODEL_NAME,
|
|
get_upstream_kwargs,
|
|
logger,
|
|
)
|
|
from schemas import DenseEmbeddingResponse, Question, SparseVector
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_sparse_model() -> SparseTextEmbedding:
|
|
logger.info("Loading local sparse model %s", SPARSE_MODEL_NAME)
|
|
return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
|
|
|
|
|
|
async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
|
|
response = await client.post(
|
|
str(EMBEDDINGS_DENSE_URL),
|
|
**get_upstream_kwargs(),
|
|
json={
|
|
"model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
|
|
"input": [text],
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
payload = DenseEmbeddingResponse.model_validate(response.json())
|
|
if not payload.data:
|
|
raise ValueError("Dense embedding response is empty")
|
|
return payload.data[0].embedding
|
|
|
|
|
|
async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
|
|
"""Single request for multiple texts — avoids N parallel calls and rate limiting."""
|
|
response = await client.post(
|
|
str(EMBEDDINGS_DENSE_URL),
|
|
**get_upstream_kwargs(),
|
|
json={
|
|
"model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
|
|
"input": texts,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
payload = DenseEmbeddingResponse.model_validate(response.json())
|
|
payload.data.sort(key=lambda x: x.index)
|
|
return [item.embedding for item in payload.data]
|
|
|
|
|
|
def embed_sparse(text: str) -> SparseVector:
|
|
vectors = list(get_sparse_model().embed([text]))
|
|
if not vectors:
|
|
raise ValueError("Sparse embedding response is empty")
|
|
item = vectors[0]
|
|
return SparseVector(
|
|
indices=[int(i) for i in item.indices.tolist()],
|
|
values=[float(v) for v in item.values.tolist()],
|
|
)
|
|
|
|
|
|
def _normalize_query(text: str) -> str:
|
|
return re.sub(r"\s+", " ", text).strip()
|
|
|
|
|
|
def build_primary_query(question: Question) -> str:
|
|
q = question.search_text.strip() if question.search_text else ""
|
|
if not q:
|
|
q = question.text.strip()
|
|
return _normalize_query(q)
|
|
|
|
|
|
def build_extra_dense_queries(question: Question) -> list[str]:
|
|
extras: list[str] = []
|
|
for v in question.variants or []:
|
|
q = _normalize_query(v)
|
|
if q:
|
|
extras.append(q)
|
|
for h in question.hyde or []:
|
|
q = _normalize_query(h)
|
|
if q:
|
|
extras.append(q)
|
|
return extras
|
|
|
|
|
|
def build_sparse_query(question: Question) -> str:
|
|
kws = question.keywords or []
|
|
if kws:
|
|
return " ".join(kws)
|
|
return build_primary_query(question)
|
|
|
|
|
|
def build_entity_tokens(question: Question) -> list[str]:
|
|
tokens: list[str] = []
|
|
if question.entities:
|
|
for field in (
|
|
question.entities.people,
|
|
question.entities.emails,
|
|
question.entities.documents,
|
|
question.entities.names,
|
|
question.entities.links,
|
|
):
|
|
tokens.extend(field or [])
|
|
return [t.strip() for t in tokens if t.strip()]
|