demosec-app / app.py
Habieb143's picture
Update app.py
b646b10 verified
raw
history blame contribute delete
No virus
1.8 kB
import streamlit as st
import re
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)
return qa
def preprocess_text(text):
# Remove special characters and punctuation
text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
# Convert to lowercase
text = text.lower()
return text
def format_answer(answer):
# Add answer formatting logic here
# For example, add bold formatting
return f"**{answer}**"
def get_answers(qa, question, text, max, min, do_sample):
try:
answers = qa(question=question, context=text, max_answer_len=max, min_answer_len=min, do_sample=do_sample)
return format_answer(answers['answer'])
except Exception as e:
st.error(f"Error: {str(e)}")
qa = load_qa_model()
st.title("Ask Questions about your Text")
sentence = st.text_area('Please paste your article :', height=30)
question = st.text_input("Questions from this article?")
button = st.button("Get me Answers")
with st.sidebar:
max = st.slider('Select max answer length', 50, 500, step=10, value=150)
min = st.slider('Select min answer length', 10, 450, step=10, value=50)
do_sample = st.checkbox("Do sample", value=False)
if button and sentence and question:
with st.spinner("Discovering Answers.."):
text = preprocess_text(sentence)
answer = get_answers(qa, question, text, max, min, do_sample)
st.write(answer)
else:
st.error("Please enter a question and text!")