main.py

import streamlit as st
import fitz  # PyMuPDF
import faiss
import numpy as np
import re
from sentence_transformers import SentenceTransformer, CrossEncoder
import ollama

# -----------------------------
# 1️⃣ PDF Text Extraction
# -----------------------------
def extract_text_from_pdf(uploaded_file):
    pdf_text = ""
    with fitz.open(stream=uploaded_file.read(), filetype="pdf") as doc:
        for page in doc:
            pdf_text += page.get_text("text")
    return pdf_text


# -----------------------------
# 2️⃣ Text Chunking
# -----------------------------
def chunk_text(text, chunk_size=500, overlap=100):
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)
    return chunks


# -----------------------------
# 3️⃣ Vector Store Creation
# -----------------------------
@st.cache_resource
def create_vector_store(chunks):
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    embeddings = embedder.encode(chunks, convert_to_numpy=True, normalize_embeddings=True)
    dim = embeddings.shape[1]
    index = faiss.IndexFlatL2(dim)
    index.add(embeddings)
    return index, chunks, embedder


# -----------------------------
# 4️⃣ Multi-Query Retrieval (Improved)
# -----------------------------
def retrieve_context_multiquery(query, embedder, index, chunks, k=3):
    reform_prompt = f"""
You are a helpful assistant that reformulates a user question into multiple semantically similar queries.
Generate 3 alternative versions that use different phrasing or synonyms.
Original question: "{query}"
Return each query on a new line.
"""
    try:
        response = ollama.generate(
            model="deepseek-r1:1.5b",
            prompt=reform_prompt,
            options={"temperature": 0.6, "max_tokens": 200}
        )
        reformulated_text = response["response"]
        alt_queries = re.findall(r'^\s*[\d\-•]?\s*(.+)', reformulated_text, re.MULTILINE)
        alt_queries = [q.strip() for q in alt_queries if q.strip()]
    except Exception as e:
        st.error(f"Query reformulation failed: {e}")
        alt_queries = []
    # Add the original question too
    all_queries = [query] + alt_queries
    st.write("🔁 Reformulated Queries:", all_queries)

    # Retrieve chunks for each query
    all_contexts = set()
    for q in all_queries:
        q_emb = embedder.encode([q], convert_to_numpy=True, normalize_embeddings=True)
        _, indices = index.search(q_emb.astype(np.float32), k)
        for i in indices[0]:
            all_contexts.add(chunks[i])

    return list(all_contexts)

# -----------------------------
# 5️⃣ Re-ranking Step
# -----------------------------
@st.cache_resource
def load_reranker():
    return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank_contexts(query, contexts, top_n=5):
    reranker = load_reranker()
    pairs = [(query, ctx) for ctx in contexts]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(contexts, scores), key=lambda x: x[1], reverse=True)
    top_contexts = [ctx for ctx, _ in ranked[:top_n]]
    return top_contexts


# -----------------------------
# 6️⃣ Answer Generation (Ollama)
# -----------------------------
def remove_think_tags(text):
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)

def generate_answer(query, context):
    context_text = "\n".join(context)
    prompt = f"""
You are a helpful AI assistant answering questions from a PDF document.

Context:
{context_text}

Question: {query}

Answer comprehensively using only the provided context:
"""
    response = ollama.generate(
        model="deepseek-r1:1.5b",
        prompt=prompt,
        options={"temperature": 0.3, "max_tokens": 1000}
    )
    return remove_think_tags(response["response"])


# -----------------------------
# 7️⃣ Streamlit UI
# -----------------------------
st.title("🧠 Multi-Query + Reranking RAG PDF Chatbot")

uploaded_file = st.file_uploader("📤 Upload a PDF file", type=["pdf"])
query = st.text_area("💬 Ask a question about your PDF:")

if uploaded_file:
    pdf_text = extract_text_from_pdf(uploaded_file)
    chunks = chunk_text(pdf_text)
    index, docs, embedder = create_vector_store(chunks)

    if st.button("🔍 Get Smart Answer"):
        with st.spinner("Thinking with Multi-Query + Reranking..."):
            contexts = retrieve_context_multiquery(query, embedder, index, docs)
            ranked_contexts = rerank_contexts(query, contexts, top_n=5)
            answer = generate_answer(query, ranked_contexts)
        st.markdown(f"### 🧠 Answer:\n{answer}")
else:
    st.info("Please upload a PDF to begin.")