revise bug in categoricalemb for graphdit

Author: liugangcode

tests/generator/graphdit.py

@@ -6,7 +6,7 @@
 from torch_molecule import GraphDITMolecularGenerator
 from torch_molecule.utils.search import ParameterType, ParameterSpec
 
-EPOCHS = 10
+EPOCHS = 2
 BATCH_SIZE = 32
 
 def test_graph_dit_generator():
@@ -18,13 +18,15 @@ def test_graph_dit_generator():
         'CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F'
     ]
     smiles_list = smiles_list * 25  # Create 100 molecules for training
-    properties = [1.0, 2.0, 3.0, 4.0] * 25  # Create 100 properties for training
+    # properties = [1.0, 2.0, 3.0, 4.0] * 25  # Create 100 properties for training
+    properties = [0, 0, 1, 1] * 25  # Create 100 properties for training
 
     # 1. Basic initialization test - Conditional Model
     print('smiles_list', len(smiles_list), smiles_list[:5], 'properties', len(properties), properties[:5])
     print("\n=== Testing Conditional GraphDIT model initialization ===")
     conditional_model = GraphDITMolecularGenerator(
-        task_type=['regression'],
+        task_type=['classification'],
+        drop_condition=0.1,
         timesteps=500,
         batch_size=BATCH_SIZE,
         epochs=EPOCHS,
@@ -42,7 +44,7 @@ def test_graph_dit_generator():
 
     # 3. Conditional generation test
     print("\n=== Testing Conditional GraphDIT generation ===")
-    target_properties = [1.0, 2.0, 3.0, 4.0]
+    target_properties = [0, 0, 1, 1]
     generated_smiles = conditional_model.generate(target_properties, batch_size=BATCH_SIZE)
     print(f"Conditionally generated {len(generated_smiles)} molecules")
     print("Example conditionally generated SMILES:", generated_smiles[:2])

torch_molecule/generator/graph_dit/modeling_graph_dit.py

@@ -468,6 +468,8 @@ def generate(self, labels: Optional[Union[List[List], np.ndarray, torch.Tensor]]
             y = labels.to(self.device).float()
         else:
             y = None
+
+        self.model.eval()
         for s_int in reversed(range(0, self.timesteps)):
             s_array = s_int * torch.ones((batch_size, 1)).float().to(self.device)
             t_array = s_array + 1

torch_molecule/nn/embedder.py

@@ -86,29 +86,6 @@ def __init__(self, num_classes, hidden_size, dropout_prob):
         self.num_classes = num_classes
         self.dropout_prob = dropout_prob
 
-    def token_drop(self, labels, force_drop_ids=None):
-        """
-        Drops labels to enable classifier-free guidance.
-
-        Parameters
-        ----------
-        labels : torch.Tensor
-            Tensor of integer labels.
-        force_drop_ids : torch.Tensor or None, optional
-            Boolean mask to force specific labels to be dropped.
-
-        Returns
-        -------
-        torch.Tensor
-            Labels with some entries replaced by a dropout token.
-        """
-        if force_drop_ids is None:
-            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
-        else:
-            drop_ids = force_drop_ids == 1
-        labels = torch.where(drop_ids, self.num_classes, labels)
-        return labels
-
     def forward(self, labels, train, force_drop_ids=None):
         """
         Forward pass for categorical embedding with optional label dropout.
@@ -128,11 +105,21 @@ def forward(self, labels, train, force_drop_ids=None):
             Embedded label representations, with optional noise added during training.
         """
         labels = labels.long().view(-1)
+
         use_dropout = self.dropout_prob > 0
-        if (train and use_dropout) or (force_drop_ids is not None):
-            labels = self.token_drop(labels, force_drop_ids)
+        drop_ids = force_drop_ids == 1
+
+        if (train and use_dropout):
+            drop_ids_rand = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+            if force_drop_ids is not None:
+                drop_ids = torch.logical_or(drop_ids, drop_ids_rand)
+            else:
+                drop_ids = drop_ids_rand
+        
+        if use_dropout:
+            labels = torch.where(drop_ids, self.num_classes, labels)
         embeddings = self.embedding_table(labels)
-        if True and train:
+        if train:
             noise = torch.randn_like(embeddings)
             embeddings = embeddings + noise
         return embeddings