import streamlit as st
from transformers import pipeline

def main():
    st.title("🤖 AxionX Digital — Question Answering Demo")

    st.markdown("Type a paragraph in **Context**, then ask a **Question** about it.")

    @st.cache_resource
    def load_model():
        return pipeline("question-answering", model="distilbert-base-cased-distilled-squad")

    qa = load_model()

    context = st.text_area("Context", value=(
        "AxionX Digital builds model-training tools for AI developers. "
        "We fine-tune open-source LLMs for customer support, finance, and legal use cases. "
        "We also provide evaluation dashboards and fast private deployments."
    ), height=200)

    question = st.text_input("Question", value="What does AxionX Digital build?")

    if st.button("🔍 Get Answer"):
        if not context.strip() or not question.strip():
            st.warning("Please enter both context and question.")
        else:
            result = qa(question=question, context=context)
            st.subheader("📄 Answer")
            st.write(f"**Answer:** {result['answer']}")
            st.write(f"**Confidence Score:** {round(result['score'], 3)}")

# ✅ Make sure it runs standalone too
if __name__ == "__main__":
    main()