|
|
| import torch, numpy as np, faiss |
| from sentence_transformers import SentenceTransformer, CrossEncoder |
|
|
| class ApexRetriever: |
| def __init__(self, model_dir="."): |
| self.bi = SentenceTransformer(f"{model_dir}/bi_encoder", device="cuda" if torch.cuda.is_available() else "cpu") |
| self.reranker = CrossEncoder(f"{model_dir}/reranker", device="cuda" if torch.cuda.is_available() else "cpu") |
| self._index, self._documents = None, [] |
|
|
| def index_documents(self, documents): |
| self._documents = documents |
| emb = self.bi.encode(documents, normalize_embeddings=True, show_progress_bar=False) |
| self._index = faiss.IndexFlatIP(emb.shape[1]) |
| self._index.add(emb.astype("float32")) |
|
|
| def retrieve(self, query, top_k=5, recall_k=100): |
| if self._index is None: raise ValueError("Index documents first.") |
| q_emb = self.bi.encode(query, normalize_embeddings=True).astype("float32") |
| _, indices = self._index.search(np.expand_dims(q_emb, 0), min(recall_k, len(self._documents))) |
| candidates = [self._documents[i] for i in indices[0]] |
| pairs = [(query, d) for d in candidates] |
| scores = self.reranker.predict(pairs) |
| return [d for _, d in sorted(zip(scores, candidates), reverse=True)[:top_k]] |
|
|
