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!")