import streamlit as st
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer

@st.cache(allow_output_mutation=True)
def load_qa_model():
    model_name = "mrm8488/mobilebert-uncased-finetuned-squadv2"
    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    qa = pipeline("question-answering", model=model, tokenizer=tokenizer)
    text_generator = pipeline("text-generation")
    return qa, text_generator

qa, text_generator = load_qa_model()
st.title("Ask Questions about your Text")
sentence = st.text_area('Please paste your article :', height=30)
num_questions = st.number_input("Number of questions to generate:", min_value=1, max_value=10, value=3, step=1)
num_answers = st.number_input("Number of answers per question:", min_value=1, max_value=5, value=1, step=1)
button = st.button("Generate Questions and Answers")

with st.spinner("Generating Questions and Answers.."):
    if button and sentence:
        generated_questions = text_generator(sentence, max_length=100, num_return_sequences=num_questions)
        for question_index, question_output in enumerate(generated_questions):
            st.subheader(f"Question {question_index + 1}: {question_output['generated_text']}")
            answers = qa(question=question_output['generated_text'], context=sentence, topk=num_answers)
            for answer_index, answer in enumerate(answers):
                st.write(f"Answer {answer_index + 1}: {answer['answer']}")