| """bge query embedding + FAISS search over the hosted foto-index.""" |
| import numpy as np |
| import pandas as pd |
| from pathlib import Path |
|
|
| BGE_PREFIX = "Represent this sentence for searching relevant passages: " |
| INDEX_REPO = "htohfa/foto-index" |
|
|
|
|
| class Retriever: |
| def __init__(self, repo=INDEX_REPO): |
| import faiss |
| from huggingface_hub import snapshot_download |
| from sentence_transformers import SentenceTransformer |
| local = Path(snapshot_download(repo, repo_type="dataset")) |
| self.index = faiss.read_index(str(next(local.glob("*.faiss")))) |
| meta = pd.read_parquet(next(local.glob("meta.parquet"))) |
| self.meta = meta.rename(columns={"fig_idx": "fig_index"}).reset_index(drop=True) |
| self.model = SentenceTransformer("BAAI/bge-base-en-v1.5") |
|
|
| def search(self, query, k=100): |
| v = self.model.encode([BGE_PREFIX + query], |
| normalize_embeddings=True).astype(np.float32) |
| sims, ids = self.index.search(v, k) |
| out = [] |
| for s, i in zip(sims[0], ids[0]): |
| if i < 0: |
| continue |
| row = self.meta.iloc[int(i)] |
| out.append({"arxiv_id": str(row.arxiv_id), |
| "fig_index": int(row.fig_index), |
| "caption_prefix": str(row.caption), |
| "score": float(s)}) |
| return out |
|
|
