Update app.py

app.py

@@ -120,12 +120,37 @@ elif page == "Data Upload":
     uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])
     
     if uploaded_file:
-        df = pd.read_csv(uploaded_file)
-        st.write("Uploaded Data Preview:", df.head())
-        
-        if st.button("Predict for Dataset"):
-            df["Prediction"] = df.apply(lambda row: predict_susceptibility(row.to_dict(), model, encoders)["Final Output"], axis=1)
-            st.write("Prediction Results", df)
+        try:
+            df = pd.read_csv(uploaded_file)
+            st.write("Uploaded Data Preview:", df.head())
+            
+            # Validate the required columns
+            required_columns = ['organism', 'antibiotic', 'was_positive']
+            missing_columns = [col for col in required_columns if col not in df.columns]
+            
+            if missing_columns:
+                st.error(f"The uploaded CSV is missing the following required columns: {', '.join(missing_columns)}")
+            else:
+                # Process predictions
+                if st.button("Predict for Dataset"):
+                    with st.spinner("Processing predictions..."):
+                        df["Prediction"] = df.apply(
+                            lambda row: predict_susceptibility(row.to_dict(), model, encoders)["Final Output"], axis=1
+                        )
+                    
+                    st.success("Predictions complete!")
+                    st.write("Prediction Results:", df)
+                    
+                    # Optionally, download the results as a CSV
+                    csv = df.to_csv(index=False)
+                    st.download_button(
+                        label="Download Results as CSV",
+                        data=csv,
+                        file_name="predictions.csv",
+                        mime='text/csv'
+                    )
+        except Exception as e:
+            st.error(f"Error reading the uploaded CSV file: {e}")
 
 # About Page
 elif page == "About":