Fix data shape mismatch

utils.py

@@ -59,7 +59,7 @@ def predict_permeability(
 
     if type(mol_features) == pd.DataFrame:
         if mol_features.index.tolist() != info_df.index.tolist():
-            raise ValueError("Features_df and Info_df do not have the same index.")
+            raise ValueError("mol_features and Info_df do not have the same index.")
 
     # get predicted probabilities
     info_df.loc[:, "B3clf_predicted_probability"] = pred_model.predict_proba(
@@ -99,7 +99,7 @@ def generate_predictions(
         if mol_features is None and info_df is None:
             if input_fname is None:
                 raise ValueError("Either input_fname or mol_features/info_df must be provided")
-                
+
             mol_tag = os.path.basename(input_fname).split(".")[0]
             file_ext = os.path.splitext(input_fname)[1].lower()
             internal_sdf = f"{mol_tag}_optimized_3d.sdf"
@@ -136,7 +136,11 @@ def generate_predictions(
                 mol_features = mol_features.apply(pd.to_numeric, errors='coerce')
                 if mol_features.isnull().any().any():
                     st.warning("Some descriptors contained invalid values and were removed")
-                    mol_features = mol_features.dropna()
+                    # Get indices of valid rows
+                    valid_indices = ~mol_features.isnull().any(axis=1)
+                    # Update both dataframes to keep only valid rows
+                    mol_features = mol_features[valid_indices]
+                    info_df = info_df[valid_indices]
                     if len(mol_features) == 0:
                         raise ValueError("No valid data remains after cleaning")
 
@@ -174,7 +178,7 @@ def generate_predictions(
         ]
 
         return mol_features, info_df, result_df
-        
+
     except Exception as e:
         import traceback
         st.error(f"Error in generate_predictions: {str(e)}\n{traceback.format_exc()}")