Spaces:
Running
Running
File size: 4,836 Bytes
c5ced83 | 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 | """Hybrid search over the openai/math catalogue.
Each family is indexed as one document (title + summary + aliases + plain-English
line) plus one document per paper (title + abstract). A family's dense score is
the max cosine similarity over its documents; its BM25 score is the max BM25
score over the same documents. The two family rankings are fused with
reciprocal-rank fusion. This is the configuration that scored best on the
eval (see eval/RESULTS.md: qwen3-0.6b+expanded, hybrid_rrf).
"""
import json
import threading
from collections import OrderedDict
from pathlib import Path
import bm25s
import numpy as np
import Stemmer
MODEL_ID = "Qwen/Qwen3-Embedding-0.6B"
# Document vectors in index/ were computed with this exact revision; queries must use it too.
MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3"
QUERY_PROMPT = ("Instruct: Given a search query about mathematics, retrieve research results "
"whose statement answers or matches the query\nQuery: ")
RRF_K = 60
def family_doc_text(f: dict) -> str:
return f"{f['title']}. {f['summary']} Also known as: {', '.join(f['aliases'])}. {f['plain']}"
def paper_doc_text(p: dict) -> str:
return f"{p['title']}. {p['abstract']}"
def build_docs(families: list[dict]) -> list[dict]:
"""Documents in family order: each family's summary doc, then its papers."""
docs = []
for fi, f in enumerate(families):
docs.append({"family": fi, "paper": None, "text": family_doc_text(f)})
for pi, p in enumerate(f["papers"]):
docs.append({"family": fi, "paper": pi, "text": paper_doc_text(p)})
return docs
class SearchEngine:
def __init__(self, index_dir: Path, load_model: bool = True):
self.catalog = json.loads((index_dir / "catalog.json").read_text())
self.families = self.catalog["families"]
self.docs = build_docs(self.families)
self.emb = np.load(index_dir / "embeddings.npy")
assert self.emb.shape[0] == len(self.docs), "index out of date: rebuild with pipeline/build_index.py"
doc_family = np.array([d["family"] for d in self.docs])
self.starts = np.flatnonzero(np.r_[True, doc_family[1:] != doc_family[:-1]])
self.ends = np.r_[self.starts[1:], len(self.docs)]
self.stemmer = Stemmer.Stemmer("english")
self.bm25 = bm25s.BM25()
self.bm25.index(bm25s.tokenize([d["text"] for d in self.docs], stopwords="en",
stemmer=self.stemmer, show_progress=False), show_progress=False)
self.model = None
self._lock = threading.Lock()
self._cache: OrderedDict[str, np.ndarray] = OrderedDict()
if load_model:
from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(MODEL_ID, revision=MODEL_REVISION, device="cpu")
self.model.max_seq_length = 512
def embed_query(self, query: str) -> np.ndarray:
if query in self._cache:
self._cache.move_to_end(query)
return self._cache[query]
with self._lock:
vec = self.model.encode([query], prompt=QUERY_PROMPT, normalize_embeddings=True,
convert_to_numpy=True)[0].astype(np.float32)
self._cache[query] = vec
if len(self._cache) > 2048:
self._cache.popitem(last=False)
return vec
def _family_max(self, doc_scores: np.ndarray) -> np.ndarray:
return np.maximum.reduceat(doc_scores, self.starts)
def search(self, query: str) -> list[dict]:
"""All families ranked for `query`, best first."""
sims = self.emb @ self.embed_query(query)
dense = self._family_max(sims)
tokens = bm25s.tokenize([query], stopwords="en", stemmer=self.stemmer,
return_ids=False, show_progress=False)[0]
lexical = self._family_max(self.bm25.get_scores(tokens)) if tokens else np.zeros_like(dense)
dense_rank = np.empty(len(dense), dtype=int)
dense_rank[np.argsort(-dense, kind="stable")] = np.arange(len(dense))
lex_rank = np.empty(len(lexical), dtype=int)
lex_rank[np.argsort(-lexical, kind="stable")] = np.arange(len(lexical))
fused = 1 / (RRF_K + dense_rank + 1) + (lexical > 0) / (RRF_K + lex_rank + 1)
results = []
for fi in np.argsort(-fused, kind="stable"):
best = self.starts[fi] + int(np.argmax(sims[self.starts[fi]:self.ends[fi]]))
results.append({
"id": self.families[fi]["id"],
"score": round(float(fused[fi]), 6),
"similarity": round(float(dense[fi]), 4),
"keyword_match": bool(lexical[fi] > 0),
"matched_paper": self.docs[best]["paper"],
})
return results
|