Spaces:
Running
Running
| """Cross-encoder reranker that fixes the order of candidates already recalled by the server. | |
| Embedding recall scores each memory on its own, so when a distractor shares the query's wording the | |
| right memory can land second: the README's weak case, "Who is presenting at the quarterly review?", | |
| ranks "The team agreed to ship the billing fix before the quarterly review." just above the memory | |
| that actually answers it. A cross-encoder reads the (query, memory) pair together and scores | |
| relevance directly, which separates those two. It runs on CPU over the handful of candidates the | |
| server already returned, so it adds one small forward pass per recall and never touches the store. | |
| """ | |
| from __future__ import annotations | |
| from collections.abc import Callable, Sequence | |
| from functools import lru_cache | |
| from typing import Any | |
| MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2" # ~80 MB, CPU, trained for query-passage relevance | |
| def _model(): | |
| # Imported lazily so the module is cheap to import and the model loads once, on first recall. | |
| from sentence_transformers import CrossEncoder | |
| return CrossEncoder(MODEL) | |
| def rerank(query: str, candidates: Sequence[Any], *, text: Callable[[Any], str] = lambda c: c.text) -> list[Any]: | |
| """Return ``candidates`` ordered by cross-encoder relevance to ``query``, highest first. | |
| The sort is stable, so candidates the reranker cannot separate keep the server's order, and a | |
| list shorter than two is returned unchanged. ``text`` reads the string to score from each | |
| candidate (``m.text`` for a recalled memory). | |
| """ | |
| items = list(candidates) | |
| if len(items) < 2: | |
| return items | |
| scores = _model().predict([(query, text(c)) for c in items]) | |
| order = sorted(range(len(items)), key=lambda i: scores[i], reverse=True) | |
| return [items[i] for i in order] | |