# Gradio App

import numpy as np
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
from sentence_transformers import SentenceTransformer
import gradio as gr
from PyPDF2 import PdfReader

# Document Reading 
def read_document(file_path):
    """
    Reads PDF or TXT files and returns plain text.
    """
    if file_path.endswith(".pdf"):
        reader = PdfReader(file_path)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + " "
        return text
    elif file_path.endswith(".txt"):
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()
    else:
        return None

# Chunking
def chunk_text(text, chunk_size=500):
    """
    Splits the text into chunks of `chunk_size` words.
    Transformers have input limits, so chunking is necessary.
    """
    words = text.split()
    chunks = [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
    return chunks
# Embeddings for retrieval
embed_model = SentenceTransformer('all-MiniLM-L6-v2')

# Pre-trained QA model for extractive answers
qa_model_name = "distilbert-base-uncased-distilled-squad"
qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
qa_model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name)
qa_pipeline_model = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
# QA Function
def answer_question(file, question, top_k=3):
    """
    Takes a document file and question string,
    returns the best answer from the document.
    """
    if file is None:
        return "Please upload a document."

    # Step 1: Read & chunk document
    text = read_document(file.name)
    chunks = chunk_text(text, chunk_size=500)

    # Step 2: Embed chunks
    chunk_embeddings = embed_model.encode(chunks)

    # Step 3: Embed question
    question_embedding = embed_model.encode([question])[0]

    # Step 4: Compute cosine similarity to find top-k relevant chunks
    similarities = np.dot(chunk_embeddings, question_embedding) / (
        np.linalg.norm(chunk_embeddings, axis=1) * np.linalg.norm(question_embedding)
    )
    top_idx = similarities.argsort()[-top_k:][::-1]
    top_chunks = [chunks[i] for i in top_idx]

    # Step 5: Run extractive QA on top-k chunks
    best_answer = {"score": 0, "answer": "Answer not found."}
    for chunk in top_chunks:
        result = qa_pipeline_model(question=question, context=chunk)
        if result['score'] > best_answer['score']:
            best_answer = result

    return best_answer['answer']

iface = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.File(label="Upload Document (.pdf or .txt)"),
        gr.Textbox(label="Enter your question"),
        gr.Slider(1, 5, value=3, step=1, label="Top-k chunks to consider")
    ],
    outputs=gr.Textbox(label="Answer"),
    title="Document Question Answering (Mini-RAG)",
    description="Upload a document and ask questions. The system retrieves relevant chunks and extracts the answer."
)

iface.launch()