73 lines
2 KiB
Python
73 lines
2 KiB
Python
import asyncio
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from config import RERANK_LIMIT, RERANKER_MODEL, RERANKER_URL, get_upstream_kwargs, logger
|
|
from retrieval import extract_page_content
|
|
|
|
|
|
async def _get_rerank_scores(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
targets: list[str],
|
|
) -> list[float]:
|
|
if not targets:
|
|
return []
|
|
|
|
for attempt in range(5):
|
|
try:
|
|
response = await client.post(
|
|
str(RERANKER_URL),
|
|
**get_upstream_kwargs(),
|
|
json={
|
|
"model": RERANKER_MODEL,
|
|
"encoding_format": "float",
|
|
"text_1": query,
|
|
"text_2": targets,
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
if attempt < 4:
|
|
await asyncio.sleep(2 ** attempt)
|
|
continue
|
|
raise exc
|
|
|
|
if response.status_code == 429:
|
|
wait = 2 ** attempt
|
|
logger.warning("Rerank 429, retry %d/5 in %ds", attempt + 1, wait)
|
|
await asyncio.sleep(wait)
|
|
continue
|
|
|
|
response.raise_for_status()
|
|
data = response.json().get("data") or []
|
|
return [float(sample["score"]) for sample in data]
|
|
|
|
logger.error("Rerank 429 after all retries, falling back")
|
|
return []
|
|
|
|
|
|
async def rerank_points(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
points: list[Any],
|
|
) -> tuple[list[Any], list[Any]]:
|
|
if not points:
|
|
return [], []
|
|
|
|
head = points[:RERANK_LIMIT]
|
|
tail = points[RERANK_LIMIT:]
|
|
targets = [extract_page_content(p) for p in head]
|
|
|
|
try:
|
|
scores = await _get_rerank_scores(client, query, targets)
|
|
except Exception as exc:
|
|
logger.warning("Rerank failed, using retrieval order: %s", exc)
|
|
return head, tail
|
|
|
|
if len(scores) != len(head):
|
|
logger.warning("Rerank score count mismatch, using retrieval order")
|
|
return head, tail
|
|
|
|
reranked = [p for _, p in sorted(zip(scores, head), key=lambda x: x[0], reverse=True)]
|
|
return reranked, tail
|