Update smiles_cor.py

smiles_cor.py

@@ -1286,6 +1286,15 @@ class smi_correct(object):
         
         df_output["SMILES"] = df_output.apply(lambda row: self.standardization_pipeline(row["CORRECT"]), axis=1)
         
-        df_output = self.remove_smiles_duplicates(df_output, subset="SMILES").drop(columns=["CORRECT", "INCORRECT", "ORIGINAL"]).dropna()
-        
+    df_output = self.remove_smiles_duplicates(df_output, subset="SMILES")
+    
+    # List of columns to drop
+    columns_to_drop = ["CORRECT", "ORIGINAL"]
+    
+    # Check if "INCORRECT" column exists and add it to the list
+    if "INCORRECT" in df_output.columns:
+        columns_to_drop.append("INCORRECT")
+    
+    # Drop the specified columns
+    df_output = df_output.drop(columns=columns_to_drop).dropna()        
         return df_output