Spaces:
Sleeping
Sleeping
| """Optional pretrained neural scorers. Weights remain with their original authors.""" | |
| from __future__ import annotations | |
| from typing import Any, Mapping, Sequence | |
| import numpy as np | |
| from .core import Passage, rank_with_scores | |
| DENSE_MODEL = "sentence-transformers/all-MiniLM-L6-v2" | |
| CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2" | |
| class EvidenceGraphReranker: | |
| """Hugging Face-friendly neural scoring + the same fixed graph rule. | |
| backend='cross-encoder' (the standard MS MARCO reranker) or backend='dense' | |
| (cosine over normalized all-MiniLM embeddings). | |
| Models are downloaded from Hugging Face at first inference. Pass scores | |
| explicitly to `rank` to avoid downloading or loading a neural model. | |
| """ | |
| def __init__(self, backend: str = "cross-encoder", device: str = "cpu", batch_size: int = 16): | |
| if backend not in {"cross-encoder", "dense"}: | |
| raise ValueError("backend must be 'cross-encoder' or 'dense'") | |
| if batch_size < 1: | |
| raise ValueError("batch_size must be positive") | |
| self.backend, self.device, self.batch_size = backend, device, batch_size | |
| self._model = None | |
| def _get_model(self): | |
| if self._model is None: | |
| try: | |
| from sentence_transformers import CrossEncoder, SentenceTransformer | |
| except ImportError as exc: | |
| raise ImportError("Install the optional dependency: pip install sentence-transformers") from exc | |
| if self.backend == "cross-encoder": | |
| self._model = CrossEncoder(CROSS_ENCODER_MODEL, max_length=256, device=self.device) | |
| else: | |
| self._model = SentenceTransformer(DENSE_MODEL, device=self.device) | |
| self._model.max_seq_length = 128 | |
| return self._model | |
| def score(self, query: str, passages: Sequence[Passage | Mapping[str, Any] | str]) -> np.ndarray: | |
| docs = [Passage.from_value(p) for p in passages] | |
| if not docs: | |
| raise ValueError("At least one passage is required") | |
| texts = [(p.title + ". " if p.title else "") + p.text for p in docs] | |
| model = self._get_model() | |
| if self.backend == "cross-encoder": | |
| out = model.predict([(query, text) for text in texts], | |
| batch_size=self.batch_size, show_progress_bar=False) | |
| return np.asarray(out, dtype=np.float32).reshape(-1) | |
| emb = model.encode([query] + texts, batch_size=self.batch_size, | |
| convert_to_numpy=True, normalize_embeddings=True, | |
| show_progress_bar=False) | |
| return (emb[1:] @ emb[0]).astype(np.float32) | |
| def rank(self, query: str, passages: Sequence[Passage | Mapping[str, Any] | str], | |
| k: int = 3, scores: Sequence[float] | None = None) -> list[dict[str, Any]]: | |
| docs = [Passage.from_value(p) for p in passages] | |
| scores = self.score(query, docs) if scores is None else scores | |
| return rank_with_scores(query, docs, scores, k=k) | |