Spaces:
Runtime error
Runtime error
| import json | |
| import faiss | |
| import numpy as np | |
| from pathlib import Path | |
| from sentence_transformers import SentenceTransformer | |
| class RAGRetriever: | |
| def __init__(self, index_dir: str, | |
| embed_model: str = "pritamdeka/S-PubMedBert-MS-MARCO"): | |
| self.index_dir = Path(index_dir) | |
| self.index = faiss.read_index(str(self.index_dir / "faiss.index")) | |
| with open(self.index_dir / "chunks.jsonl") as f: | |
| self.chunks = [json.loads(line) for line in f] | |
| self.encoder = SentenceTransformer(embed_model) | |
| def search(self, query: str, k: int = 4, min_score: float = 0.3) -> list[dict]: | |
| q = self.encoder.encode([query], normalize_embeddings=True) | |
| scores, idxs = self.index.search(np.asarray(q, dtype="float32"), k) | |
| return [ | |
| {**self.chunks[i], "score": float(scores[0][j])} | |
| for j, i in enumerate(idxs[0]) | |
| if i >= 0 and scores[0][j] >= min_score | |
| ] |