Update streamlit_app.py

streamlit_app.py

@@ -2,20 +2,27 @@ import streamlit as st
 import pandas as pd
 from transformers import pipeline
 
-
-st.set_page_config(page_title="News Classifier", page_icon="📰")
+st.set_page_config(page_title="News Classifier + Q&A", page_icon="📰")
 
 @st.cache_resource
-def load_model():
+def load_classifier():
     return pipeline(
         "text-classification",
         model="SuchithTissera/text_analytics_model"
     )
 
-classifier = load_model()
+@st.cache_resource
+def load_qa_model():
+    return pipeline(
+        "question-answering",
+        model="distilbert-base-cased-distilled-squad"
+    )
+
+classifier = load_classifier()
+qa_model = load_qa_model()
 
-st.title("News Classifier")
-st.write("Upload a CSV file with a `content` column. The app will add a new `class` column.")
+st.title("News Classifier + Q&A")
+st.write("Upload a CSV file with a `content` column. The app will classify each record, add a `class` column, and let you ask questions about a selected article.")
 
 uploaded_file = st.file_uploader("Upload CSV", type=["csv"])
 
@@ -34,6 +41,8 @@ if uploaded_file is not None:
             df["class"] = [pred["label"] for pred in preds]
 
             st.success("Classification complete.")
+
+            st.subheader("Preview of Classified Data")
             st.dataframe(df.head())
 
             output_csv = df.to_csv(index=False).encode("utf-8")
@@ -45,5 +54,39 @@ if uploaded_file is not None:
                 mime="text/csv"
             )
 
+            st.markdown("---")
+            st.subheader("Question Answering")
+
+            article_options = [
+                f"Article {i}: {text[:100]}..." if len(text) > 100 else f"Article {i}: {text}"
+                for i, text in enumerate(df["content"].fillna("").astype(str).tolist())
+            ]
+
+            selected_index = st.selectbox(
+                "Choose an article",
+                options=range(len(article_options)),
+                format_func=lambda i: article_options[i]
+            )
+
+            selected_article = df["content"].fillna("").astype(str).iloc[selected_index]
+
+            st.write("### Selected Article")
+            st.write(selected_article)
+
+            question = st.text_input("Ask a question about the selected article")
+
+            if st.button("Get Answer"):
+                if not question.strip():
+                    st.warning("Please enter a question.")
+                elif not selected_article.strip():
+                    st.warning("The selected article is empty.")
+                else:
+                    with st.spinner("Finding answer..."):
+                        result = qa_model(question=question, context=selected_article)
+
+                    st.success("Answer found.")
+                    st.write("### Answer")
+                    st.write(result["answer"])
+
     except Exception as e:
         st.error(f"Error reading or processing file: {e}")