Update app.py

app.py

@@ -176,11 +176,12 @@ elif page == "Data Upload":
                         progress_bar = st.progress(0)
                         total_rows = len(df)
                         
-                        # Initialize results columns
-                        df["Prediction"] = ""
-                        df["Rule Guidance"] = ""
-                        df["Model Prediction"] = ""
-                        df["Decision Reason"] = ""
+                        # Create a new results DataFrame with the same index as the original
+                        results_df = pd.DataFrame(index=df.index)
+                        results_df["Prediction"] = ""
+                        results_df["Rule Guidance"] = ""
+                        results_df["Model Prediction"] = ""
+                        results_df["Decision Reason"] = ""
                         
                         # Process each row with error handling
                         for i, (index, row) in enumerate(df.iterrows()):
@@ -188,25 +189,35 @@ elif page == "Data Upload":
                                 # Skip rows with invalid data
                                 if (row['organism'] not in encoders['organism'] or 
                                     row['antibiotic'] not in encoders['antibiotic']):
-                                    df.at[index, "Prediction"] = "Invalid data"
+                                    results_df.at[index, "Prediction"] = "Invalid data"
                                     continue
                                 
+                                # Extract only the required columns for prediction
+                                input_data = {
+                                    'organism': row['organism'],
+                                    'antibiotic': row['antibiotic'],
+                                    'was_positive': row['was_positive']
+                                }
+                                
                                 # Get full prediction result
-                                result = predict_susceptibility(row.to_dict(), model, encoders)
+                                result = predict_susceptibility(input_data, model, encoders)
                                 
                                 # Store all results
                                 if "Error" in result:
-                                    df.at[index, "Prediction"] = "Error: " + result["Error"]
+                                    results_df.at[index, "Prediction"] = "Error: " + result["Error"]
                                 else:
-                                    df.at[index, "Prediction"] = result["Final Output"]
-                                    df.at[index, "Rule Guidance"] = result["Rule Guidance"]
-                                    df.at[index, "Model Prediction"] = result["Model Prediction"]
-                                    df.at[index, "Decision Reason"] = result["Decision Reason"]
+                                    results_df.at[index, "Prediction"] = result["Final Output"]
+                                    results_df.at[index, "Rule Guidance"] = result["Rule Guidance"]
+                                    results_df.at[index, "Model Prediction"] = result["Model Prediction"]
+                                    results_df.at[index, "Decision Reason"] = result["Decision Reason"]
                             except Exception as e:
-                                df.at[index, "Prediction"] = f"Error: {str(e)}"
+                                results_df.at[index, "Prediction"] = f"Error: {str(e)}"
                             
                             # Update progress bar
                             progress_bar.progress((i + 1) / total_rows)
+                        
+                        # Combine original data with results
+                        df = pd.concat([df, results_df], axis=1)
                     
                     st.success("Predictions complete!")