| from __future__ import annotations |
|
|
| import math |
|
|
| from .domain import ( |
| CandidateScorer, |
| CandidateScoringError, |
| RankedCandidate, |
| RerankConfig, |
| RerankRequest, |
| RerankResult, |
| ) |
|
|
|
|
| class Reranker: |
| def __init__(self, scorer: CandidateScorer, config: RerankConfig | None = None) -> None: |
| self._scorer = scorer |
| self._config = config or RerankConfig() |
|
|
| def rerank(self, request: RerankRequest) -> RerankResult: |
| if len(request.candidates) < 2: |
| return self._preserve(request, reason="insufficient_candidates") |
| if len(request.candidates) > self._config.max_candidates: |
| return self._preserve(request, reason="unsafe_input") |
| if any( |
| len(candidate.surface) > self._config.max_candidate_chars |
| for candidate in request.candidates |
| ): |
| return self._preserve(request, reason="unsafe_input") |
| surfaces = [candidate.surface for candidate in request.candidates] |
| if ( |
| any(not surface for surface in surfaces) |
| or len(set(surfaces)) != len(surfaces) |
| or any(not math.isfinite(candidate.prior_score) for candidate in request.candidates) |
| ): |
| return self._preserve(request, reason="unsafe_input") |
| if ( |
| len(request.left_context) + len(request.right_context) |
| > self._config.max_context_words |
| ): |
| return self._preserve(request, reason="unsafe_input") |
| if len(request.reading) > self._config.max_reading_chars: |
| return self._preserve(request, reason="unsafe_input") |
| if ( |
| sum(len(word) for word in request.left_context + request.right_context) |
| > self._config.max_context_chars |
| ): |
| return self._preserve(request, reason="unsafe_input") |
| try: |
| model_scores = tuple( |
| float(score) for score in self._scorer.score_candidates(request) |
| ) |
| except CandidateScoringError as error: |
| return self._preserve(request, reason=error.code) |
| except Exception: |
| return self._preserve(request, reason="scorer_error") |
| if len(model_scores) != len(request.candidates) or not all( |
| math.isfinite(score) for score in model_scores |
| ): |
| return self._preserve(request, reason="invalid_scores") |
| combined_scores = tuple( |
| model_score + self._config.prior_weight * candidate.prior_score |
| for candidate, model_score in zip(request.candidates, model_scores, strict=True) |
| ) |
| model_order = sorted( |
| range(len(request.candidates)), |
| key=lambda index: (-combined_scores[index], index), |
| ) |
| margin = combined_scores[model_order[0]] - combined_scores[model_order[1]] |
| changed = model_order[0] != 0 and margin >= self._config.min_margin |
| order = model_order if changed else list(range(len(request.candidates))) |
| if changed: |
| reason = "accepted" |
| elif model_order[0] == 0: |
| reason = "baseline_best" |
| else: |
| reason = "low_margin" |
| ranked = tuple( |
| RankedCandidate( |
| surface=request.candidates[index].surface, |
| original_rank=index, |
| prior_score=request.candidates[index].prior_score, |
| model_score=model_scores[index], |
| combined_score=combined_scores[index], |
| ) |
| for index in order |
| ) |
| return RerankResult( |
| reading=request.reading, |
| ranked=ranked, |
| changed=changed, |
| reason=reason, |
| margin=margin, |
| ) |
|
|
| @staticmethod |
| def _preserve(request: RerankRequest, *, reason: str) -> RerankResult: |
| return RerankResult( |
| reading=request.reading, |
| ranked=tuple( |
| RankedCandidate( |
| surface=candidate.surface, |
| original_rank=index, |
| prior_score=candidate.prior_score, |
| model_score=None, |
| combined_score=None, |
| ) |
| for index, candidate in enumerate(request.candidates) |
| ), |
| changed=False, |
| reason=reason, |
| margin=None, |
| ) |
|
|
