import gradio as gr
import pdfplumber
from transformers import pipeline

# Load HuggingFace free models
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
qa_model = pipeline("question-answering", model="google/flan-t5-small")

def extract_text_from_pdf(pdf_file):
    text = ""
    with pdfplumber.open(pdf_file.name) as pdf:
        for page in pdf.pages:
            text += page.extract_text() + "\n"
    return text

def summarize_pdf(pdf_file):
    text = extract_text_from_pdf(pdf_file)
    summary = summarizer(text, max_length=200, min_length=50, do_sample=False)
    return summary[0]['summary_text']

def generate_qa(pdf_file, question):
    text = extract_text_from_pdf(pdf_file)
    answer = qa_model(question=question, context=text)
    return answer['answer']

with gr.Blocks() as demo:
    gr.Markdown("# PDF Summarizer + Q&A (Free)")

    with gr.Tab("Summarize PDF"):
        pdf_input = gr.File(label="Upload PDF")
        summary_output = gr.Textbox(label="Summary", lines=10)
        summarize_btn = gr.Button("Generate Summary")
        summarize_btn.click(summarize_pdf, inputs=pdf_input, outputs=summary_output)

    with gr.Tab("PDF Q&A"):
        pdf_input_qa = gr.File(label="Upload PDF")
        question_input = gr.Textbox(label="Ask a Question")
        answer_output = gr.Textbox(label="Answer", lines=5)
        qa_btn = gr.Button("Get Answer")
        qa_btn.click(generate_qa, inputs=[pdf_input_qa, question_input], outputs=answer_output)

demo.launch()