forked from zovos/vk_hackathon
e244b8bf30
- embed_dense_multi now sends one batch request (N texts → 1 API call) instead of N parallel requests, avoiding rate-limit errors when question has variants/hyde - Extra dense embeddings (variants/hyde) wrapped in try/except so primary query always succeeds - Reranker now retries up to 5 times with exponential backoff on 429, matching Lotus reference Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
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 []
|
|
|
|
import asyncio as _asyncio
|
|
|
|
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,
|
|
},
|
|
)
|
|
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]
|
|
except Exception as exc:
|
|
if attempt < 4:
|
|
await _asyncio.sleep(2 ** attempt)
|
|
continue
|
|
raise exc
|
|
|
|
return []
|
|
|
|
|
|
async def rerank_points(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
points: list[Any],
|
|
) -> tuple[list[Any], list[Any]]:
|
|
"""Return (reranked_head, retrieval_tail) so we don't lose candidates."""
|
|
if not points:
|
|
return [], []
|
|
|
|
rerank_candidates = points[:RERANK_LIMIT]
|
|
tail = points[RERANK_LIMIT:]
|
|
|
|
targets = [extract_page_content(p) for p in rerank_candidates]
|
|
try:
|
|
scores = await get_rerank_scores(client, query, targets)
|
|
except Exception as exc:
|
|
logger.warning("Rerank failed, using retrieval order: %s", exc)
|
|
return rerank_candidates, tail
|
|
|
|
if len(scores) != len(rerank_candidates):
|
|
logger.warning("Rerank score count mismatch, using retrieval order")
|
|
return rerank_candidates, tail
|
|
|
|
paired = sorted(zip(scores, rerank_candidates), key=lambda x: x[0], reverse=True)
|
|
reranked = [p for _, p in paired]
|
|
return reranked, tail
|