Update app.py

app.py

@@ -1,54 +1,34 @@
-import os
-import gradio as gr
+import streamlit as st
 from transformers import pipeline
 
-# --- Prevent storage overflow: redirect caches to /tmp (auto-cleared) ---
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
-os.environ["HF_HOME"] = "/tmp/hf_home"
-os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_datasets"
-
-# --- Pin model version for reproducibility ---
-MODEL_ID = "distilbert-base-cased-distilled-squad"
-
-# Lazy-load the model only once at startup
-def load_model():
-    return pipeline("question-answering", model=MODEL_ID)
-
-qa = load_model()
-
-# --- Demo defaults ---
-EXAMPLE_CONTEXT = (
-    "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."
-)
-EXAMPLE_QUESTION = "What does AxionX Digital build?"
-
-# --- Prediction function ---
-def predict(context, question):
-    context = (context or "").strip()
-    question = (question or "").strip()
-    if not context or not question:
-        return {"answer": "Please provide both context and a question.", "score": ""}
-    res = qa(question=question, context=context)
-    return {"answer": res.get("answer", ""), "score": round(float(res.get("score", 0.0)), 3)}
-
-# --- Gradio UI (with AxionX brand style) ---
-with gr.Blocks(title="AxionX Digital — AI QA Demo") as demo:
-    gr.Markdown(
-        """
-        # 🤖 **AxionX Digital — Question Answering Demo**
-        Type a paragraph in **Context**, then ask a **Question** about it.  
-        *(Model: distilbert-base-cased-distilled-squad — lightweight for stability)*  
-        ---
-        """
-    )
-    with gr.Row():
-        ctx = gr.Textbox(label="Context", lines=8, value=EXAMPLE_CONTEXT)
-    q = gr.Textbox(label="Question", value=EXAMPLE_QUESTION)
-    btn = gr.Button("🔍 Get Answer")
-    out = gr.JSON(label="Result (answer, score)")
-    btn.click(predict, inputs=[ctx, q], outputs=[out])
+def main():
+    st.set_page_config(page_title="Gradio QA Demo", page_icon="🤖", layout="wide")
+    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)}")
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    main()