Spaces:
Running on Zero
Running on Zero
File size: 3,512 Bytes
465b477 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Within-document retrievers used by the RAG baseline and the retrieval diagnostics."""
from __future__ import annotations
import re
from functools import lru_cache
import numpy as np
from rank_bm25 import BM25Okapi
from .schemas import Chunk
_TOKEN = re.compile(r"[a-z0-9]+")
_STOP = frozenset(
[
"a",
"an",
"the",
"of",
"to",
"in",
"on",
"for",
"and",
"or",
"is",
"are",
"was",
"were",
"be",
"been",
"by",
"with",
"as",
"at",
"from",
"that",
"this",
"these",
"those",
"it",
"its",
"which",
"what",
"who",
"whom",
"how",
"why",
"when",
"where",
"do",
"does",
"did",
"their",
"they",
"them",
"we",
"our",
"you",
"your",
"not",
"no",
]
)
def tokenize(text: str) -> list[str]:
return [t for t in _TOKEN.findall(text.lower()) if t not in _STOP]
def bm25_rank(question: str, chunks: list[Chunk]) -> list[tuple[int, float]]:
"""Indices of ``chunks`` sorted by BM25 score (ties broken by document order)."""
if not chunks:
return []
bm25 = BM25Okapi([tokenize(c.section_heading + " " + c.text) or ["_"] for c in chunks])
scores = bm25.get_scores(tokenize(question))
return sorted(((i, float(s)) for i, s in enumerate(scores)), key=lambda t: (-t[1], t[0]))
@lru_cache(maxsize=2)
def _encoder(model_name: str):
from sentence_transformers import SentenceTransformer
return SentenceTransformer(model_name, device="cpu")
class DenseRanker:
"""Cosine-similarity ranking with a sentence-transformers bi-encoder (optional extra ``nli``)."""
def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
self.model_name = model_name
self._cache: dict[str, np.ndarray] = {}
def _embed(self, texts: list[str], keys: list[str]) -> np.ndarray:
missing = [(k, t) for k, t in zip(keys, texts) if k not in self._cache]
if missing:
vecs = _encoder(self.model_name).encode([t for _, t in missing], normalize_embeddings=True, batch_size=16)
for (k, _), v in zip(missing, vecs):
self._cache[k] = v
return np.stack([self._cache[k] for k in keys])
def rank(self, question: str, chunks: list[Chunk]) -> list[tuple[int, float]]:
if not chunks:
return []
q = _encoder(self.model_name).encode([question], normalize_embeddings=True)[0]
m = self._embed([c.section_heading + ". " + c.text for c in chunks], [c.chunk_id for c in chunks])
scores = m @ q
return sorted(((i, float(s)) for i, s in enumerate(scores)), key=lambda t: (-t[1], t[0]))
def fill_budget(order: list[int], costs: list[int], budget_tokens: int, max_items: int | None = None) -> list[int]:
"""Take items in ``order`` while their token ``costs`` fit the budget (skipping ones that do not).
Costs must be measured on the exact text the model will see, so every system is held
to the same budget.
"""
picked, used = [], 0
for i in order:
if max_items is not None and len(picked) >= max_items:
break
t = costs[i]
if used + t <= budget_tokens:
picked.append(i)
used += t
return picked
|