Update app.py

app.py

@@ -1,80 +1,75 @@
 import streamlit as st
 import pandas as pd
-import time
-from transformers import pipeline
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+import io
 
-# Set page title and layout
-st.set_page_config(page_title="CSV News Classifier", layout="wide")
+# Load pre-trained model and tokenizer for text classification
+model_name = "TAgroup5/news-classification-model"
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Load the fine-tuned Hugging Face model
-@st.cache_resource
-def load_model():
-    return pipeline("text-classification", model="TAgroup5/daily-mirror-news-classifier")
+# Initialize the text classification pipeline
+text_classification_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
-classifier = load_model()
+# Initialize the question answering pipeline
+qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
 
-# Custom CSS for Colors and Styling
-st.markdown(
-    """
-    <style>
-        body { background-color: #f8f9fa; }
-        .stApp { background-color: #ffffff; padding: 20px; border-radius: 10px; box-shadow: 0px 0px 10px rgba(0,0,0,0.1); }
-        
-        .custom-box {
-            background-color: #f0f0f0; 
-            padding: 15px;
-            border-radius: 10px;
-            border: 1px solid #ccc;
-            margin-bottom: 15px;
-        }
-        h1 { color: #ff5733; text-align: center; }  
-        h2, h3 { color: #007bff; }  
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
-# Page Title
-st.title("📊 News Classification from CSV")
-
-# File Uploader
-st.markdown('<div class="custom-box"><h3>📂 Upload a CSV file</h3></div>', unsafe_allow_html=True)
-uploaded_file = st.file_uploader("", type=["csv"])
+# Streamlit App Layout
+st.title("News Classification and Q&A")
+
+# Component 1: Text Classification Pipeline
+st.header("Classify News Articles")
+
+st.markdown("""
+    Upload a CSV file containing news articles, and the model will classify each article 
+    into one of the following categories: Business, Opinion, Political Gossip, Sports, or World News.
+""")
+
+uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
 
 if uploaded_file is not None:
-    # Read CSV
     df = pd.read_csv(uploaded_file)
+    
+    if 'content' not in df.columns:
+        st.error("The uploaded CSV file must have a 'content' column containing news excerpts.")
+    else:
+        st.write("Preview of the data:")
+        st.dataframe(df.head())
 
-    # Show Preview
-    st.markdown("### 🔍 **Preview of Uploaded File**", unsafe_allow_html=True)
-    st.dataframe(df.head())
+        # Preprocess the data and classify each article
+        def preprocess_text(text):
+            # Apply necessary preprocessing steps here (e.g., removing stopwords, special characters, etc.)
+            return text
+        
+        # Apply preprocessing and classification
+        df['processed_content'] = df['content'].apply(preprocess_text)
+        df['class'] = df['processed_content'].apply(lambda x: text_classification_pipeline(x)[0]['label'])
+        
+        # Show the results
+        st.write("Classification Results:")
+        st.dataframe(df[['content', 'class']])
 
-    # Assuming the column with news articles is named "news_text"
-    if "news_text" not in df.columns:
-        st.error("❌ The uploaded CSV must contain a 'news_text' column.")
-    else:
-        # Perform classification
-        st.markdown("### 🏷️ **Classifying News Articles...**")
-        with st.spinner("Processing..."):
-            df["class"] = df["news_text"].apply(lambda text: classifier(text)[0]["label"])
+        # Provide an option to download the output as CSV
+        output = io.StringIO()
+        df.to_csv(output, index=False)
+        st.download_button(label="Download classified news", data=output.getvalue(), file_name="output.csv", mime="text/csv")
 
-        # Show Preview of Results
-        st.markdown("### 📌 **Preview of Classified Data**", unsafe_allow_html=True)
-        st.dataframe(df.head())
 
-        # Download Button for Classified CSV
-        csv = df.to_csv(index=False).encode("utf-8")
-        st.markdown("### 📥 **Download Classified CSV**", unsafe_allow_html=True)
-        st.download_button(
-            label="⬇️ **Download Classified CSV**",
-            data=csv,
-            file_name="classified_news.csv",
-            mime="text/csv",
-            help="Click to download the classified news file"
-        )
-
-# Footer
-st.markdown("---")
-st.markdown('<p style="text-align:center; font-size:14px; color:#6c757d;">👨‍💻 Developed by <b>Ridmi Navodya</b> | Powered by Streamlit 🚀</p>', unsafe_allow_html=True)
+# Component 2: Q&A Pipeline
+st.header("Ask a Question About the News")
 
+st.markdown("""
+    Type in a question, and the model will extract an answer from the provided news content.
+""")
+
+question = st.text_input("Ask a question:")
+
+if question:
+    context = st.text_area("Provide the news article or content for the Q&A:", height=150)
+    
+    if context:
+        # Perform the question-answering task
+        result = qa_pipeline(question=question, context=context)
+        
+        st.write("Answer:", result['answer'])