Update src/data/dataset.py

src/data/dataset.py

@@ -89,11 +89,11 @@ class DruggenDataset(InMemoryDataset):
             smiles_list (list): List of SMILES strings.
 
         Returns:
-            max_length (int): Maximum number of atoms found in the filtered molecules.
+            num_smiles (int): Number of filtered smiles
             filtered_smiles (list): List of valid SMILES strings.
         """
-        max_length = 0
         filtered_smiles = []
+        num_smiles = 0
         for smiles in tqdm(smiles_list, desc="Filtering SMILES"):
             mol = Chem.MolFromSmiles(smiles)
             if mol is None:
@@ -113,8 +113,9 @@ class DruggenDataset(InMemoryDataset):
                 continue
 
             filtered_smiles.append(smiles)
-            max_length = max(max_length, molecule_size)
-        return max_length, filtered_smiles
+            num_smiles += 1
+
+        return num_smiles, filtered_smiles
 
     def _genA(self, mol, connected=True, max_length=None):
         """
@@ -290,20 +291,22 @@ class DruggenDataset(InMemoryDataset):
         """
         # Read raw SMILES from file (assuming CSV with no header)
         smiles_list = pd.read_csv(self.raw_files, header=None)[0].tolist()
-        max_length, filtered_smiles = self._filter_smiles(smiles_list)
+        num_smiles, filtered_smiles = self._filter_smiles(smiles_list)
+        self.num_smiles = num_smiles
+
         data_list = []
         self.m_dim = len(self.atom_decoder_m)
         for smiles in tqdm(filtered_smiles, desc='Processing dataset', total=len(filtered_smiles)):
             mol = Chem.MolFromSmiles(smiles)
-            A = self._genA(mol, connected=True, max_length=max_length)
+            A = self._genA(mol, connected=True, max_length=self.max_atom)
             if A is not None:
-                x_array = self._genX(mol, max_length=max_length)
+                x_array = self._genX(mol, max_length=self.max_atom)
                 if x_array is None:
                     continue
                 x = torch.from_numpy(x_array).to(torch.long).view(1, -1)
                 x = label2onehot(x, self.m_dim).squeeze()
                 if self.features:
-                    f = torch.from_numpy(self._genF(mol, max_length=max_length)).to(torch.long).view(x.shape[0], -1)
+                    f = torch.from_numpy(self._genF(mol, max_length=self.max_atom)).to(torch.long).view(x.shape[0], -1)
                     x = torch.concat((x, f), dim=-1)
                 adjacency = torch.from_numpy(A)
                 edge_index = adjacency.nonzero(as_tuple=False).t().contiguous()