import streamlit as st
from transformers import pipeline

from transformers import AutoTokenizer, AutoModelForQuestionAnswering

model_name = "wedo2910/qa_arabic_model"
tokenizer = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabertv02")
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

qa_pipeline = pipeline(
    "question-answering",
    model=model,
    tokenizer=tokenizer
)

# Default settings
default_settings = {
    "max_new_tokens": 512,
    "temperature": 0.7,
    "top_p": 0.9,
    "min_p": 0,
    "top_k": 0,
    "repetition_penalty": 1.0,
    "presence_penalty": 0,
    "frequency_penalty": 0,
    "max_answer_len": 50,
    "doc_stride": 128,
}

# Define a default context (e.g., a general knowledge text or topic)
default_context = """
التزم بنص السؤال.
"""

# Streamlit UI
st.title("Arabic AI Question Answering")
st.subheader("Ask a question to get an answer.")

# Input field for the question only
question = st.text_input("Question", placeholder="Enter your question here...")

# Settings sliders
st.subheader("Settings")
max_new_tokens = st.number_input("Max New Tokens", min_value=1, max_value=1000000, value=512)
temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7, step=0.1)
top_p = st.slider("Top P", min_value=0.0, max_value=1.0, value=0.9, step=0.1)
min_p = st.slider("Min P", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
top_k = st.number_input("Top K", min_value=0, max_value=1000, value=0)
repetition_penalty = st.slider("Repetition Penalty", min_value=0.01, max_value=5.0, value=1.0, step=0.1)
presence_penalty = st.slider("Presence Penalty", min_value=-2.0, max_value=2.0, value=0.0, step=0.1)
frequency_penalty = st.slider("Frequency Penalty", min_value=-2.0, max_value=2.0, value=0.0, step=0.1)
max_answer_len = st.number_input("Max Answer Length", min_value=1, value=50)
doc_stride = st.number_input("Document Stride", min_value=1, value=128)

# Generate Answer button
if st.button("Get Answer"):
    if not question:
        st.error("The question field is required.")
    else:
        # Generate answer using the default context
        try:
            prediction = qa_pipeline(
                {"context": default_context, "question": question},
                max_answer_len=max_answer_len,
                doc_stride=doc_stride
            )
            st.subheader("Result")
            st.write(f"**Question:** {question}")
            st.write(f"**Answer:** {prediction['answer']}")
        except Exception as e:
            st.error(f"Error: {e}")