195 lines
7.0 KiB
Python
195 lines
7.0 KiB
Python
import asyncio
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Iterable, Sequence
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
import numpy as np
|
|
|
|
from yuxi.services.model_cache import model_cache
|
|
from yuxi.utils import get_docker_safe_url, logger
|
|
|
|
|
|
def sigmoid(x):
|
|
return 1 / (1 + np.exp(-x))
|
|
|
|
|
|
class BaseReranker(ABC):
|
|
def __init__(self, model_name, api_key, base_url, **kwargs):
|
|
self.url = get_docker_safe_url(base_url)
|
|
self.model = model_name
|
|
self.api_key = api_key
|
|
self.headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
|
|
self.session: aiohttp.ClientSession | None = None
|
|
self.timeout = aiohttp.ClientTimeout(total=30)
|
|
self.parameters: dict[str, Any] = dict(kwargs.get("parameters", {}))
|
|
|
|
async def _ensure_session(self) -> None:
|
|
if self.session is None or self.session.closed:
|
|
self.session = aiohttp.ClientSession(headers=self.headers, timeout=self.timeout)
|
|
|
|
@abstractmethod
|
|
def _build_payload(self, query: str, documents: list[str], max_length: int) -> dict[str, Any]:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def _extract_results(self, result: dict[str, Any]) -> list[dict[str, Any]]:
|
|
raise NotImplementedError
|
|
|
|
async def acompute_score(
|
|
self,
|
|
sentence_pairs: Sequence[Sequence[str]],
|
|
batch_size: int = 32,
|
|
max_length: int = 512,
|
|
normalize: bool = True,
|
|
) -> list[float]:
|
|
if not sentence_pairs or len(sentence_pairs) < 2:
|
|
return []
|
|
|
|
query, sentences = sentence_pairs[0], sentence_pairs[1]
|
|
documents = [sentences] if isinstance(sentences, str) else list(sentences)
|
|
|
|
if not documents:
|
|
return []
|
|
|
|
await self._ensure_session()
|
|
|
|
all_scores: list[float] = []
|
|
batch_size = max(1, int(batch_size))
|
|
total_batches = (len(documents) + batch_size - 1) // batch_size
|
|
|
|
for batch_no, start in enumerate(range(0, len(documents), batch_size), start=1):
|
|
batch = documents[start : start + batch_size]
|
|
try:
|
|
scores = await self._batch_rerank(query, batch, max_length=max_length)
|
|
all_scores.extend(scores)
|
|
logger.debug(f"Reranking batch {batch_no}/{total_batches} completed")
|
|
except Exception as exc:
|
|
logger.error(f"Reranking batch {batch_no} failed: {exc}")
|
|
all_scores.extend([0.5] * len(batch))
|
|
|
|
if normalize:
|
|
all_scores = [float(sigmoid(score)) for score in all_scores]
|
|
|
|
return all_scores
|
|
|
|
async def _batch_rerank(self, query: str, documents: Iterable[str], max_length: int) -> list[float]:
|
|
docs = list(documents)
|
|
if not docs:
|
|
return []
|
|
|
|
payload = self._build_payload(query, docs, max_length)
|
|
|
|
await self._ensure_session()
|
|
assert self.session is not None
|
|
|
|
try:
|
|
async with self.session.post(self.url, json=payload) as response:
|
|
response.raise_for_status()
|
|
result: dict[str, Any] = await response.json()
|
|
except TimeoutError as exc:
|
|
total_timeout = self.timeout.total if self.timeout else 0.0
|
|
logger.error(f"Reranking request timeout after {total_timeout:.1f}s")
|
|
raise exc
|
|
except aiohttp.ClientError as exc:
|
|
logger.error(f"Reranking request failed: {exc}")
|
|
raise exc
|
|
|
|
processed = sorted(self._extract_results(result), key=lambda item: item.get("index", 0))
|
|
return [float(entry.get("relevance_score", 0.0)) for entry in processed]
|
|
|
|
def compute_score(self, sentence_pairs, batch_size=256, max_length=512, normalize=False):
|
|
try:
|
|
_ = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return asyncio.run(self.acompute_score(sentence_pairs, batch_size, max_length, normalize))
|
|
raise RuntimeError("compute_score cannot be used while an event loop is running. Use acompute_score instead.")
|
|
|
|
async def test_connection(self) -> tuple[bool, str]:
|
|
try:
|
|
scores = await self._batch_rerank("test query", ["test document"], max_length=128)
|
|
if scores:
|
|
return True, "连接正常"
|
|
return False, "响应无效"
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.error(f"Rerank connection test failed: {error_msg}")
|
|
return False, error_msg
|
|
finally:
|
|
await self.aclose()
|
|
|
|
async def aclose(self) -> None:
|
|
if self.session and not self.session.closed:
|
|
await self.session.close()
|
|
|
|
def __del__(self) -> None:
|
|
if self.session and not self.session.closed:
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
asyncio.run(self.aclose())
|
|
return
|
|
|
|
if loop.is_closed():
|
|
asyncio.run(self.aclose())
|
|
elif not loop.is_running():
|
|
loop.run_until_complete(self.aclose())
|
|
|
|
|
|
class OpenAIReranker(BaseReranker):
|
|
def _build_payload(self, query: str, documents: list[str], max_length: int) -> dict[str, Any]:
|
|
return {
|
|
"model": self.model,
|
|
"query": query,
|
|
"documents": documents,
|
|
"max_chunks_per_doc": max_length,
|
|
}
|
|
|
|
def _extract_results(self, result: dict[str, Any]) -> list[dict[str, Any]]:
|
|
return list(result.get("results", []))
|
|
|
|
|
|
class DashscopeReranker(BaseReranker):
|
|
def _build_payload(self, query: str, documents: list[str], max_length: int) -> dict[str, Any]:
|
|
params = {"top_n": len(documents), "return_documents": False}
|
|
instruct = self.parameters.get("instruct")
|
|
if instruct:
|
|
params["instruct"] = instruct
|
|
return {
|
|
"model": self.model,
|
|
"input": {"query": query, "documents": documents},
|
|
"parameters": params,
|
|
}
|
|
|
|
def _extract_results(self, result: dict[str, Any]) -> list[dict[str, Any]]:
|
|
return list(result.get("output", {}).get("results", []))
|
|
|
|
|
|
def get_reranker(model_id: str, **kwargs):
|
|
info = model_cache.get_model_info(model_id)
|
|
if not info:
|
|
raise ValueError(f"Unknown reranker model spec: {model_id}")
|
|
if info.model_type != "rerank":
|
|
raise ValueError(f"Model {model_id} is not a rerank model (type={info.model_type})")
|
|
if not info.api_key:
|
|
raise ValueError(f"{info.display_name} api_key is required")
|
|
|
|
parameters = dict(info.extra.get("parameters") or {})
|
|
parameters.update(kwargs.pop("parameters", {}) or {})
|
|
reranker_kwargs = {**kwargs, "parameters": parameters}
|
|
protocol = info.extra.get("rerank_protocol") or info.extra.get("protocol")
|
|
|
|
if protocol == "dashscope":
|
|
return DashscopeReranker(
|
|
model_name=info.model_id,
|
|
api_key=info.api_key,
|
|
base_url=info.base_url,
|
|
**reranker_kwargs,
|
|
)
|
|
return OpenAIReranker(
|
|
model_name=info.model_id,
|
|
api_key=info.api_key,
|
|
base_url=info.base_url,
|
|
**reranker_kwargs,
|
|
)
|