demosec-app / app.py
Habieb143's picture
Update app.py
116a19e verified
raw
history blame
1.53 kB
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=50, 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']}")