Update pages/1_Classification.py

pages/1_Classification.py

@@ -14,11 +14,6 @@ def load_classifier():
 classifier = load_classifier()
 
 st.title("News Classifier")
-st.write(
-    "Upload a CSV file with a `content` column. The app will classify each record, "
-    "add a `class` column, and let you view the class distribution."
-)
-
 uploaded_file = st.file_uploader("Upload CSV", type=["csv"])
 
 if uploaded_file is not None:
@@ -30,37 +25,17 @@ if uploaded_file is not None:
         else:
             df["content"] = df["content"].fillna("").astype(str)
 
-            texts = df["content"].tolist()
-
             with st.spinner("Classifying records..."):
-                preds = classifier(texts, batch_size=8)
+                preds = classifier(df["content"].tolist(), batch_size=8)
 
             df["class"] = [pred["label"] for pred in preds]
 
-            st.success("Classification complete.")
+            # save for other pages
+            st.session_state["df"] = df
+            st.session_state["uploaded_filename"] = uploaded_file.name
 
-            st.subheader("Preview of Classified Data")
+            st.success("Classification complete.")
             st.dataframe(df.head())
 
-            st.subheader("Class Distribution")
-            class_counts = df["class"].value_counts()
-
-            fig, ax = plt.subplots()
-            ax.bar(class_counts.index, class_counts.values)
-            ax.set_xlabel("Class")
-            ax.set_ylabel("Number of Articles")
-            ax.set_title("Distribution of Predicted Labels")
-            plt.xticks(rotation=45)
-            st.pyplot(fig)
-
-            output_csv = df.to_csv(index=False).encode("utf-8")
-
-            st.download_button(
-                label="Download output.csv",
-                data=output_csv,
-                file_name="output.csv",
-                mime="text/csv"
-            )
-
     except Exception as e:
         st.error(f"Error reading or processing file: {e}")