import streamlit as st
from tempfile import NamedTemporaryFile
import os
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

# Function to save the uploaded PDF to a temporary file
def save_uploaded_file(uploaded_file):
    with NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        temp_file.write(uploaded_file.read())
        return temp_file.name

# Streamlit UI
st.title("PDF Question Answering App")
uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
if uploaded_file is not None:
    # Save the uploaded file to a temporary location
    temp_file_path = save_uploaded_file(uploaded_file)

    # Load the pre-trained question-answering model and tokenizer
    model_name = "deepset/roberta-base-squad2"
    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Define a function to get answers
    def get_answer(question, context):
        inputs = tokenizer(question, context, return_tensors="pt")
        start_scores, end_scores = model(**inputs)
        answer_start = start_scores.argmax()
        answer_end = end_scores.argmax() + 1
        answer = tokenizer.decode(inputs["input_ids"][0][answer_start:answer_end])
        return answer

    question = st.text_input("Enter your question:")
    context = st.text_area("Enter the context:")
    if st.button("Get Answer"):
        if not context:
            st.warning("Please provide context for the question.")
        else:
            answer = get_answer(question, context)
            st.write("Answer:")
            st.write(answer)

    # Cleanup: Delete the temporary file
    os.remove(temp_file_path)