Update app.py

app.py

@@ -47,7 +47,6 @@ def process_uploaded_file(uploaded_file, required_columns):
 
         # Standardize column names to uppercase and strip spaces
         df.columns = df.columns.str.upper().str.strip()
-        # st.write("DataFrame columns:", df.columns.tolist())
 
         # Standardize required columns to uppercase and strip spaces
         required_columns_upper = [col.upper().strip() for col in required_columns]
@@ -321,16 +320,44 @@ elif interface == "Batch Prediction":
 
         st.success('Batch predictions completed for all selected models.')
 
-        # Allow user to download the results for each model
-        for model_name, output_df in model_outputs.items():
-            csv = output_df.to_csv(index=False).encode('utf-8')
-            st.download_button(
-                label=f"Download {model_name} Predictions as CSV",
-                data=csv,
-                file_name=f'{model_name.lower().replace(" ", "_")}_predictions.csv',
-                mime='text/csv',
-            )
-
+        # Download options
+        st.header('Download Predictions')
+        download_option = st.radio(
+            "Choose how to download your predictions:",
+            ("All Models in Separate Files", "Churn and Non-Churn in Separate Files")
+        )
+
+        if download_option == "All Models in Separate Files":
+            # Allow user to download the results for each model
+            for model_name, output_df in model_outputs.items():
+                csv = output_df.to_csv(index=False).encode('utf-8')
+                st.download_button(
+                    label=f"Download {model_name} Predictions as CSV",
+                    data=csv,
+                    file_name=f'{model_name.lower().replace(" ", "_")}_predictions.csv',
+                    mime='text/csv',
+                )
+        elif download_option == "Churn and Non-Churn in Separate Files":
+            # Consolidate results for all models and split into churn and non-churn files
+            for model_name, output_df in model_outputs.items():
+                churn_df = output_df[output_df['Churn'] == 'Yes']
+                non_churn_df = output_df[output_df['Churn'] == 'No']
+
+                churn_csv = churn_df.to_csv(index=False).encode('utf-8')
+                non_churn_csv = non_churn_df.to_csv(index=False).encode('utf-8')
+
+                st.download_button(
+                    label=f"Download {model_name} Churn Predictions as CSV",
+                    data=churn_csv,
+                    file_name=f'{model_name.lower().replace(" ", "_")}_churn_predictions.csv',
+                    mime='text/csv',
+                )
+                st.download_button(
+                    label=f"Download {model_name} Non-Churn Predictions as CSV",
+                    data=non_churn_csv,
+                    file_name=f'{model_name.lower().replace(" ", "_")}_non_churn_predictions.csv',
+                    mime='text/csv',
+                )
     else:
         st.info('Awaiting CSV or Excel file to be uploaded.')
 
@@ -340,4 +367,4 @@ st.sidebar.write(f"Total models available: {len(models)}")
 st.sidebar.write(f"Models selected for prediction: {len(selected_models)}")
 st.sidebar.write("### Model Accuracies")
 for model, accuracy in model_accuracies.items():
-    st.sidebar.write(f"{model}: {accuracy}%")
+    st.sidebar.write(f"{model}: {accuracy}%")