Merge pull request #136 from ChEB-AI/fix/duplicate_smiles_aug

Avoid generation of original SMILES in augmentation

Author: sfluegel05

chebai/preprocessing/datasets/chebi.py

@@ -358,7 +358,8 @@ def _perform_smiles_augmentation(self) -> None:
             self.processed_main_file_names_dict["data"]
         )
 
-        AUG_SMILES_VARIATIONS = self.aug_smiles_variations
+        # +1 to account for if original SMILES is generated by random chance
+        AUG_SMILES_VARIATIONS = self.aug_smiles_variations + 1
 
         def generate_augmented_smiles(smiles: str) -> list[str]:
             mol: Chem.Mol = Chem.MolFromSmiles(smiles)
@@ -397,6 +398,16 @@ def generate_augmented_smiles(smiles: str) -> list[str]:
             else:
                 augmented = frag_smiles[0]
 
+            if smiles in augmented:
+                augmented.remove(smiles)
+
+            if len(augmented) > AUG_SMILES_VARIATIONS - 1:
+                # AUG_SMILES_VARIATIONS = number of new smiles needed to generate + original smiles
+                # if 3 smiles variations are needed, and 4 are generated because none
+                # correponds to original smiles and no-one is elimnated in previous if condition
+                augmented = random.sample(augmented, AUG_SMILES_VARIATIONS - 1)
+
+            # original smiles always first in the list
             return [smiles] + list(augmented)
 
         data_df["SMILES"] = data_df["SMILES"].apply(generate_augmented_smiles)