import gradio as gr
from transformers import pipeline
import pdfplumber

# Load the pre-trained question-answering model
qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")

# Shared variable to store uploaded PDF text
pdf_text = ""

# Function to load the PDF and store its text
def load_pdf(file):
    global pdf_text
    try:
        with pdfplumber.open(file) as pdf:
            pdf_text = ""
            for page in pdf.pages:
                pdf_text += page.extract_text()
        return "PDF loaded successfully."
    except Exception as e:
        return f"Error processing PDF: {str(e)}"

# Function to answer the user's question based on the loaded PDF
def answer_question(question):
    if not pdf_text:
        return "No PDF loaded. Upload a PDF first."
    
    try:
        # Ask the user's question using the question-answering model
        answer = qa_pipeline({"context": pdf_text, "question": question})
        return answer["answer"]
    except Exception as e:
        return f"Error answering question: {str(e)}"

# Interface for uploading the PDF
pdf_interface = gr.Interface(
    fn=load_pdf,
    inputs=gr.File(label="Upload PDF"),
    outputs="text",
    live=True,
    title="PDF Uploader",
    description="Upload a PDF to load its content.",
)

# Interface for answering questions based on the loaded PDF
qa_interface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(label="Enter Question", type="text"),
    outputs="text",
    live=True,
    title="PDF Question-Answering",
    description="Enter a question to get an answer based on the loaded PDF.",
)

pdf_interface.launch()
qa_interface.launch()