avoid generation of original smiles in augmentation

Author: aditya0by0

chebai/preprocessing/datasets/chebi.py

@@ -144,9 +144,9 @@ def __init__(
         **kwargs,
     ):
         if bool(augment_smiles):
-            assert (
-                int(aug_smiles_variations) > 0
-            ), "Number of variations must be greater than 0"
+            assert int(aug_smiles_variations) > 0, (
+                "Number of variations must be greater than 0"
+            )
             aug_smiles_variations = int(aug_smiles_variations)
 
             if not kwargs.get("splits_file_path", None):
@@ -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)