arsam
53c5aa9d61
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
import asyncio
|
|
import logging
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_model = None
|
|
_predict_lock = asyncio.Lock()
|
|
|
|
|
|
def load_reranker() -> None:
|
|
global _model
|
|
from axon.config import settings
|
|
from sentence_transformers import CrossEncoder
|
|
|
|
logger.info("Loading reranker model: %s", settings.reranker_model)
|
|
_model = CrossEncoder(settings.reranker_model)
|
|
logger.info("Reranker model loaded")
|
|
|
|
|
|
def get_reranker() -> "CrossEncoder | None":
|
|
return _model
|
|
|
|
|
|
def chunk_text(text: str, max_tokens: int = 512, max_chunks: int = 4) -> list[str]:
|
|
if not text:
|
|
return [""]
|
|
words = text.split()
|
|
words_per_chunk = max(1, int(max_tokens * 0.75))
|
|
chunks = [
|
|
" ".join(words[i : i + words_per_chunk])
|
|
for i in range(0, len(words), words_per_chunk)
|
|
]
|
|
return (chunks or [""])[:max_chunks]
|
|
|
|
|
|
def _sigmoid(x: np.ndarray) -> np.ndarray:
|
|
return 1.0 / (1.0 + np.exp(-x))
|
|
|
|
|
|
async def rerank_results(
|
|
query: str,
|
|
contents: list[tuple[str, str]],
|
|
) -> list[tuple[str, str, float]]:
|
|
if not contents:
|
|
return []
|
|
|
|
# Build (url, chunks) pairs and flat passage list
|
|
url_chunks: list[tuple[str, list[str]]] = []
|
|
flat_passages: list[str] = []
|
|
for url, text in contents:
|
|
chunks = chunk_text(text)
|
|
url_chunks.append((url, chunks))
|
|
flat_passages.extend(chunks)
|
|
|
|
pairs = [[query, p] for p in flat_passages]
|
|
|
|
def _predict():
|
|
return _model.predict(pairs)
|
|
|
|
async with _predict_lock:
|
|
raw_scores = await asyncio.to_thread(_predict)
|
|
scores: list[float] = _sigmoid(np.array(raw_scores, dtype=float)).tolist()
|
|
|
|
# Pick best chunk per URL
|
|
results: list[tuple[str, str, float]] = []
|
|
idx = 0
|
|
for url, chunks in url_chunks:
|
|
chunk_scores = scores[idx : idx + len(chunks)]
|
|
best_i = int(np.argmax(chunk_scores))
|
|
results.append((url, chunks[best_i], round(chunk_scores[best_i], 6)))
|
|
idx += len(chunks)
|
|
|
|
return sorted(results, key=lambda t: t[2], reverse=True)
|