| |
|
| | import os, re, json, pickle, logging, numpy as np, faiss
|
| | from tqdm.notebook import tqdm
|
| | from sentence_transformers import SentenceTransformer
|
| | from langchain_community.retrievers import BM25Retriever
|
| | from langchain.docstore.document import Document
|
| |
|
| | logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
| | logger = logging.getLogger(__name__)
|
| |
|
| | WORK = "context"
|
| | JSONL = f"{WORK}/rag_documents.jsonl"
|
| | FAISS_INDEX = f"{WORK}/faiss_ivf.index"
|
| | BM25_PICKLE = f"{WORK}/bm25_retriever.pkl"
|
| |
|
| | logger.info("Loading all RAG documents...")
|
| | with open(JSONL, encoding='utf-8') as f:
|
| | ALL_DOCS = [json.loads(line) for line in f]
|
| |
|
| | LINE_TO_TEXT = {i: doc["text"] for i, doc in enumerate(ALL_DOCS)}
|
| | LINE_TO_META = {i: doc["metadata"] for i, doc in enumerate(ALL_DOCS)}
|
| |
|
| | class HybridRetriever:
|
| | def __init__(self):
|
| |
|
| | self.faiss_index = faiss.read_index(FAISS_INDEX)
|
| | logger.info(f"FAISS loaded ({self.faiss_index.ntotal:,} vectors)")
|
| |
|
| |
|
| | self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2",
|
| | device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu")
|
| |
|
| |
|
| | if os.path.exists(BM25_PICKLE):
|
| | self.bm25 = pickle.load(open(BM25_PICKLE, "rb"))
|
| | logger.info("BM25 loaded")
|
| | else:
|
| | logger.info("Building BM25...")
|
| | docs = [Document(page_content=re.sub(r"^Filename:.*\nFullPath:.*\n\n", "",
|
| | doc["text"], flags=re.M),
|
| | metadata=doc["metadata"]) for doc in ALL_DOCS]
|
| | self.bm25 = BM25Retriever.from_documents(docs)
|
| | self.bm25.k = 30
|
| | pickle.dump(self.bm25, open(BM25_PICKLE, "wb"))
|
| | logger.info("BM25 built and saved")
|
| |
|
| | def batch_retrieve(self, queries, top_k=3, faiss_k=10, bm25_k=3):
|
| | qvecs = self.model.encode(queries, show_progress_bar=False, normalize_embeddings=True).astype("float32")
|
| | D, I = self.faiss_index.search(qvecs, faiss_k)
|
| |
|
| | batch_results = []
|
| | for qi, (scores, indices) in enumerate(zip(D, I)):
|
| | results = []
|
| | seen = set()
|
| | for score, idx in zip(scores, indices):
|
| | if idx == -1 or idx in seen: continue
|
| | results.append({"score": float(score), "text": LINE_TO_TEXT[idx],
|
| | "metadata": LINE_TO_META[idx], "source": "FAISS"})
|
| | seen.add(idx)
|
| | if len(results) >= top_k: break
|
| |
|
| |
|
| | bm25_docs = self.bm25.invoke(queries[qi])
|
| | for doc in bm25_docs[:bm25_k]:
|
| | ln = doc.metadata.get("line_no")
|
| | if ln in seen: continue
|
| | results.append({"score": 0.0, "text": LINE_TO_TEXT.get(ln, ""),
|
| | "metadata": LINE_TO_META.get(ln, doc.metadata), "source": "BM25"})
|
| | seen.add(ln)
|
| | if len(results) >= top_k: break
|
| | batch_results.append(results)
|
| | return batch_results
|
| |
|
| |
|
| | retriever = HybridRetriever()
|
| |
|
