Fix: skip single atom molecule in diffusion model data preprocessing. Fixes #14

Author: liugangcode

examples/generation/digress_on_qm9.py

@@ -0,0 +1,27 @@
+import os
+
+from torch_molecule import DigressMolecularGenerator
+from torch_molecule.datasets import load_qm9
+
+
+def train_on_qm9() -> None:
+    model = DigressMolecularGenerator(verbose=True, batch_size=1024, epochs=2)
+
+    smiles_list, _ = load_qm9(local_dir="torchmol_data")
+
+    original_count = len(smiles_list)
+    smiles_list = [s for s in smiles_list if isinstance(s, str) and s]
+    if original_count > len(smiles_list):
+        print(f"Data cleaning: removed {original_count - len(smiles_list)} invalid entries from QM9 dataset.")
+
+    model.fit(smiles_list)
+
+    print("\n=== Generating 10 molecules from QM9-trained model ===")
+    generated_smiles = model.generate(batch_size=10)
+    print(f"Generated {len(generated_smiles)} molecules.")
+    for i, smiles in enumerate(generated_smiles, start=1):
+        print(f"{i}: {smiles}")
+
+
+if __name__ == "__main__":
+    train_on_qm9()

torch_molecule/generator/digress/modeling_digress.py

@@ -167,6 +167,8 @@ def _convert_to_pytorch_data(self, X, y=None):
 
             # No H, first heavy atom has type 0
             node_type = torch.from_numpy(graph['node_feat'][:, 0] - 1)
+            if node_type.numel() <= 1:
+                continue
 
             # Filter out invalid node types (< 0)
             valid_mask = node_type >= 0
@@ -191,16 +193,16 @@ def _convert_to_pytorch_data(self, X, y=None):
                 # Update node and edge data
                 node_type = node_type[valid_mask]
                 g.edge_index = remapped_edge_index
-                g.edge_attr = valid_edge_attr.long().squeeze(-1)
+                g.edge_attr = valid_edge_attr.long()
             else:
                 # No invalid nodes, proceed normally
                 g.edge_index = torch.from_numpy(graph["edge_index"])
                 edge_attr = torch.from_numpy(graph["edge_feat"])[:, 0] + 1
-                g.edge_attr = edge_attr.long().squeeze(-1)
-            
+                g.edge_attr = edge_attr.long()
+
             # * is encoded as "misc" which is 119 - 1 and should be 117
             node_type[node_type == 118] = 117
-            g.x = node_type.long().squeeze(-1)
+            g.x = node_type.long()
             # g.y = torch.from_numpy(graph["y"])
             g.y = torch.zeros(1, 1)
             del graph["node_feat"]
@@ -273,7 +275,7 @@ def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
 
     def fit(self, X_train: List[str]) -> "DigressMolecularGenerator":
         num_task = 0 if self.input_dim_y is None else self.input_dim_y
-        X_train, _ = self._validate_inputs(X_train, num_task=num_task)
+        X_train, _ = self._validate_inputs(X_train, num_task=num_task, return_rdkit_mol=False)
         self._setup_diffusion_params(X_train)
         self._initialize_model(self.model_class)
         self.model.initialize_parameters()

torch_molecule/generator/gdss/modeling_gdss.py

@@ -235,6 +235,8 @@ def _convert_to_pytorch_data(self, X, y=None):
 
             # No H, first heavy atom has type 0
             node_type = torch.from_numpy(graph['node_feat'][:, 0] - 1)
+            if node_type.numel() <= 1:
+                continue
 
             # Filter out invalid node types (< 0)
             valid_mask = node_type >= 0

torch_molecule/generator/graph_dit/modeling_graph_dit.py

@@ -190,6 +190,8 @@ def _convert_to_pytorch_data(self, X, y=None):
 
             # No H, first heavy atom has type 0
             node_type = torch.from_numpy(graph['node_feat'][:, 0] - 1)
+            if node_type.numel() <= 1:
+                continue
 
             # Filter out invalid node types (< 0)
             valid_mask = node_type >= 0