rag.py

import os
import PyPDF2
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from rapidfuzz import fuzz

embedder = SentenceTransformer("all-MiniLM-L6-v2")
faiss_index = None
pdf_chunks = []
chunk_texts = []

def process_pdfs(pdf_files):
    global faiss_index, pdf_chunks, chunk_texts
    all_text = ""
    chunk_texts = []
    for pdf_file in pdf_files:
        reader = PyPDF2.PdfReader(pdf_file.name)
        for page in reader.pages:
            all_text += page.extract_text() + "\n"
    chunk_size = 500
    pdf_chunks = [all_text[i:i+chunk_size] for i in range(0, len(all_text), chunk_size)]
    chunk_texts = pdf_chunks
    embeddings = embedder.encode(pdf_chunks, convert_to_numpy=True)
    dim = embeddings.shape[1]
    faiss_index = faiss.IndexFlatL2(dim)
    faiss_index.add(embeddings)
    return f"Processed {len(pdf_chunks)} chunks from {len(pdf_files)} PDF(s)."

def semantic_search(query, top_k=3):
    global faiss_index, chunk_texts
    if faiss_index is None or not chunk_texts:
        return []
    query_emb = embedder.encode([query], convert_to_numpy=True)
    D, I = faiss_index.search(query_emb, top_k)
    return [chunk_texts[i] for i in I[0] if i < len(chunk_texts)]

def keyword_search(query, top_k=3):
    global chunk_texts
    if not chunk_texts:
        return []
    scored = [(chunk, fuzz.partial_ratio(query.lower(), chunk.lower())) for chunk in chunk_texts]
    scored = sorted(scored, key=lambda x: x[1], reverse=True)
    return [chunk for chunk, score in scored[:top_k]]

def retrieve_context(query, top_k=3):
    semantic_results = semantic_search(query, top_k)
    keyword_results = keyword_search(query, top_k)
    combined = []
    seen = set()
    for chunk in semantic_results + keyword_results:
        if chunk not in seen:
            combined.append(chunk)
            seen.add(chunk)
    return "\n".join(combined)