File size: 3,739 Bytes
c223b53 376909e c223b53 376909e c223b53 376909e c223b53 376909e c223b53 376909e c223b53 376909e c223b53 376909e c223b53 874f8b6 c223b53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """
RAG Engine for Sakhi.
Fetches top chunks from FAISS, deduplicates, and merges contiguous text blocks.
"""
import os
import pickle
import logging
from config import EMBED_MODEL, TOP_K_FETCH, TOP_K_MERGED, FAISS_INDEX_PATH, CHUNKS_PATH
logger = logging.getLogger(__name__)
class RAGEngine:
def __init__(self):
self.model = None
self.index = None
self.chunks = []
self._initialized = False
def initialize(self):
if self._initialized:
return
from build_index import ensure_index_exists
ensure_index_exists()
try:
import faiss
import torch
from sentence_transformers import SentenceTransformer
device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = SentenceTransformer(
EMBED_MODEL,
device=device,
)
self.index = faiss.read_index(FAISS_INDEX_PATH)
with open(CHUNKS_PATH, "rb") as f:
self.chunks = pickle.load(f)
logger.info(
f"Loaded {len(self.chunks)} chunks."
)
self._initialized = True
except Exception as e:
logger.exception(e)
def retrieve(self, query: str) -> list[dict]:
self.initialize()
if not self.chunks or self.index is None or self.model is None:
return []
try:
import faiss
# Fetch wider net of chunks initially
query_embedding = self.model.encode([query], convert_to_numpy=True, normalize_embeddings=True)
scores, indices = self.index.search(query_embedding, min(TOP_K_FETCH, len(self.chunks)))
retrieved = []
for score, idx in zip(scores[0], indices[0]):
if 0 <= idx < len(self.chunks):
chunk = self.chunks[idx].copy()
chunk["score"] = float(score)
retrieved.append(chunk)
# Intelligent Merging: Sort by source, then by chunk_idx
# This ensures adjacent paragraphs in the PDF are merged into one readable block
unique_chunks = {c['chunk_idx']: c for c in retrieved}.values()
sorted_chunks = sorted(unique_chunks, key=lambda x: (x.get('source', ''), x.get('chunk_idx', 0)))
merged_results = []
curr = None
for c in sorted_chunks:
if not curr:
curr = c.copy()
else:
# If chunks are from the same document and are adjacent/overlapping
if curr['source'] == c['source'] and (c['chunk_idx'] - curr['chunk_idx']) <= 2:
curr['text'] += "\n" + c['text']
curr['chunk_idx'] = c['chunk_idx'] # Move the end pointer
else:
merged_results.append(curr)
curr = c.copy()
if curr:
merged_results.append(curr)
# Return top merged blocks
return sorted(merged_results, key=lambda x: x.get('score', 0), reverse=True)[:TOP_K_MERGED]
except Exception as e:
logger.error(f"FAISS retrieval error: {e}")
return []
def format_context(self, chunks: list[dict]) -> str:
if not chunks:
return "No specific PDF context found. Please use general knowledge to explain."
context_parts = []
for i, chunk in enumerate(chunks, 1):
source = chunk.get("source", "Unknown")
context_parts.append(f"--- Document: {source} ---\n{chunk.get('text', '')}\n")
return "\n".join(context_parts) |