Spaces:
Running
Running
| """Vector store wrapper around Chroma for relevance scoring of papers. | |
| The embedding model is loaded once per process (singleton) and computed | |
| embeddings are cached on disk by content hash, so repeated papers across runs | |
| skip re-encoding. Each run still uses its own Chroma collection so concurrent | |
| jobs never contaminate one another's candidate set. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import os | |
| import uuid | |
| from .tools import Paper | |
| EMBED_MODEL = "all-MiniLM-L6-v2" | |
| CACHE_DIR = os.getenv("EMBED_CACHE_DIR", "papers/.embcache") | |
| _model = None # lazily-loaded SentenceTransformer singleton | |
| def _get_model(): | |
| global _model | |
| if _model is None: | |
| from sentence_transformers import SentenceTransformer | |
| _model = SentenceTransformer(EMBED_MODEL) | |
| return _model | |
| def _embed(texts: list[str]) -> list[list[float]]: | |
| """Embed texts, using an on-disk per-text cache keyed by content hash.""" | |
| import numpy as np | |
| os.makedirs(CACHE_DIR, exist_ok=True) | |
| out: list[list[float] | None] = [None] * len(texts) | |
| to_compute: list[int] = [] | |
| paths: list[str] = [] | |
| for i, t in enumerate(texts): | |
| key = hashlib.sha1(t.encode("utf-8")).hexdigest() | |
| path = os.path.join(CACHE_DIR, f"{key}.npy") | |
| paths.append(path) | |
| if os.path.exists(path): | |
| try: | |
| out[i] = np.load(path).tolist() | |
| continue | |
| except Exception: | |
| pass | |
| to_compute.append(i) | |
| if to_compute: | |
| vecs = _get_model().encode([texts[i] for i in to_compute]) | |
| for j, i in enumerate(to_compute): | |
| vec = np.asarray(vecs[j], dtype="float32") | |
| out[i] = vec.tolist() | |
| try: | |
| np.save(paths[i], vec) | |
| except Exception: | |
| pass | |
| return [v for v in out] # type: ignore[return-value] | |
| class PaperVectorDB: | |
| """In-memory Chroma collection (per run) over cached MiniLM embeddings.""" | |
| def __init__(self): | |
| import chromadb | |
| self._client = chromadb.Client() | |
| self._name = f"papers_{uuid.uuid4().hex[:8]}" | |
| self._collection = self._client.get_or_create_collection(name=self._name) | |
| def _document(paper: Paper) -> str: | |
| first_claim = paper.claims[0] if paper.claims else "" | |
| return f"{paper.title}\n{paper.abstract}\n{first_claim}".strip() | |
| def index(self, papers: list[Paper]) -> None: | |
| if not papers: | |
| return | |
| docs = [self._document(p) for p in papers] | |
| self._collection.upsert( | |
| ids=[p.id for p in papers], | |
| documents=docs, | |
| embeddings=_embed(docs), | |
| metadatas=[{"title": p.title, "year": p.year} for p in papers], | |
| ) | |
| def search(self, query: str, n: int = 10) -> list[str]: | |
| """Return paper IDs ranked by similarity to ``query`` (best first).""" | |
| count = self._collection.count() | |
| if count == 0: | |
| return [] | |
| res = self._collection.query( | |
| query_embeddings=_embed([query]), n_results=min(n, count) | |
| ) | |
| ids = res.get("ids") or [[]] | |
| return ids[0] | |
| def close(self) -> None: | |
| """Drop this run's collection to free memory.""" | |
| try: | |
| self._client.delete_collection(self._name) | |
| except Exception: | |
| pass | |