import streamlit as st
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
from pathlib import Path
from pdfminer.high_level import extract_text

def main():
    st.title("QUESTION - ANSWERING")
    # Upload the pdf
    file = st.file_uploader("Upload a PDF file", type=["pdf"])
    
    if file is not None:
        # Extract text from pdf
        text = extract_text(file)

        # Get model and tokenizer
        model_name = "deepset/roberta-base-squad2"
        model = AutoModelForQuestionAnswering.from_pretrained(model_name)
        tokenizer = AutoTokenizer.from_pretrained(model_name)

        # Get predictions
        nlp = pipeline('question-answering', model=model, tokenizer=tokenizer)
        
        # Get user's question
        question = st.text_input("Enter your question")

        # Check if the question is not empty
        if question:
            QA_input = {
                'question': question,
                'context': text
            }
            res = nlp(QA_input)

            # Display the answer
            st.write(res['answer'])
        else:
            # Display message when the question is empty
            st.write("Please, ask your question")

if __name__ == "__main__":
    main()