Spaces:
Running
Running
| """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 | |