forked from zovos/vk_hackathon
e49e3417eb
- Add logviewer/: Dozzle web UI (port 9999) + analyze.py CLI tool - docker-compose.yml: add json-file logging with rotation and labels for index/search - Fix Dockerfiles: COPY *.py . so all modules are included in image - Convert all relative imports to flat absolute imports for Docker flat layout - Rename index/schemas.py → index/index_schemas.py to avoid module name collision with search/schemas.py in test runner - Update all tests to add service dir to sys.path and use flat imports Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
58 lines
1.6 KiB
Python
58 lines
1.6 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 []
|
|
|
|
response = await client.post(
|
|
str(RERANKER_URL),
|
|
**get_upstream_kwargs(),
|
|
json={
|
|
"model": RERANKER_MODEL,
|
|
"encoding_format": "float",
|
|
"text_1": query,
|
|
"text_2": targets,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
|
|
data = response.json().get("data") or []
|
|
return [float(sample["score"]) for sample in data]
|
|
|
|
|
|
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
|