import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForQuestionAnswering
from sentence_transformers import SentenceTransformer
import fitz  # PyMuPDF
import os
import torch  # Import torch module

# Load the models
summarization_model_name = 'facebook/bart-large-cnn'
tokenizer = AutoTokenizer.from_pretrained(summarization_model_name)
summarization_model = AutoModelForSeq2SeqLM.from_pretrained(summarization_model_name)

qa_model_name = 'deepset/bert-large-uncased-whole-word-masking-squad2'
qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
qa_model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name)

# Function to extract text from a PDF file
def extract_text_from_pdf(file):
    doc = fitz.open(file)
    text = ""
    for page in doc:
        text += page.get_text()
    return text

# Function to summarize document
def summarize_document(document):
    inputs = tokenizer(document, return_tensors='pt', max_length=1024, truncation=True)
    summary_ids = summarization_model.generate(inputs['input_ids'], max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)

# Function to get answer to question
def get_answer(question, context):
    inputs = qa_tokenizer(question, context, return_tensors="pt", padding=True, truncation=True, max_length=512)
    outputs = qa_model(**inputs)
    start_scores = outputs.start_logits
    end_scores = outputs.end_logits

    answer_start = torch.argmax(start_scores)
    answer_end = torch.argmax(end_scores) + 1
    answer = qa_tokenizer.convert_tokens_to_string(qa_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
    return answer

# Streamlit app
st.title("PDF Summarizer and Q&A")
st.write("Upload a PDF file to get a summary and ask questions about the content.")

uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")

if uploaded_file is not None:
    # Save the uploaded file to a temporary location
    with open("temp.pdf", "wb") as f:
        f.write(uploaded_file.getbuffer())
    
    # Extract text from the PDF
    document_text = extract_text_from_pdf("temp.pdf")
    
    # Display the extracted text
    st.write("Extracted Text:")
    st.write(document_text)
    
    if st.button("Summarize"):
        with st.spinner('Summarizing...'):
            summary = summarize_document(document_text)
            st.write("**Summary:**")
            st.write(summary)
    
    question = st.text_input("Ask a question about the document")
    
    if st.button("Get Answer"):
        if question:
            with st.spinner('Generating answer...'):
                answer = get_answer(question, document_text)
                st.write("**Answer:**")
                st.write(answer)
        else:
            st.write("Please enter a question.")

# Remove temporary file after use
if os.path.exists("temp.pdf"):
    os.remove("temp.pdf")