Update PolyFusion/GINE.py

PolyFusion/GINE.py

@@ -1,11 +1,6 @@
 """
 GINE.py
-GINE-based masked pretraining on polymer graphs.
-
-This file provides (and uses internally):
-- safe_get, match_edge_attr_to_index  (used by CL.py)
-- GineEncoder (used by CL.py AND used internally by MaskedGINE)
-- MaskedGINE training script unchanged in behavior (still predicts masked atoms + hop anchors)
+GINE-based masked pretraining on polymer 2D graphs.
 """
 
 from __future__ import annotations
@@ -287,7 +282,7 @@ def compute_class_weights(train_atomic: List[torch.Tensor]) -> torch.Tensor:
 
 
 # =============================================================================
-# Encoder wrapper used by CL.py AND used internally by MaskedGINE
+# Encoder wrapper used by MaskedGINE
 # =============================================================================
 
 class GineBlock(nn.Module):
@@ -309,7 +304,7 @@ class GineBlock(nn.Module):
 
 class GineEncoder(nn.Module):
     """
-    Graph encoder for CL.py:
+    Graph encoder:
     - Produces node embeddings via GINE
     - Provides pooled graph embedding via mean pooling + pool_proj
     - Provides node_logits(...) for reconstruction (atomic prediction head)
@@ -338,10 +333,10 @@ class GineEncoder(nn.Module):
 
         self.gnn_layers = nn.ModuleList([GineBlock(node_emb_dim) for _ in range(num_layers)])
 
-        # node head for masked-atom reconstruction (used in CL.py)
+        # node head for masked-atom reconstruction 
         self.atom_head = nn.Linear(node_emb_dim, MASK_ATOM_ID + 1)
 
-        # pooled embedding projection for CL.py
+        # pooled embedding projection 
         self.pool_proj = nn.Linear(node_emb_dim, emb_dim)
 
         if class_weights is not None:
@@ -376,7 +371,7 @@ class GineEncoder(nn.Module):
 
     def forward(self, z, chirality, formal_charge, edge_index, edge_attr, batch=None):
         """
-        Returns pooled graph embedding (B, emb_dim) for CL.py.
+        Returns pooled graph embedding (B, emb_dim).
         Pool = mean over nodes per graph (batch vector).
         """
         if batch is None:
@@ -401,7 +396,7 @@ class GineEncoder(nn.Module):
 
 
 # =============================================================================
-# Training dataset + collate (unchanged behavior)
+# Training dataset + collate
 # =============================================================================
 
 class PolymerDataset(Dataset):
@@ -554,7 +549,7 @@ def collate_batch(batch):
 
 
 # =============================================================================
-# Masked pretraining model (uses GineEncoder internally)
+# Masked pretraining model
 # =============================================================================
 
 class MaskedGINE(nn.Module):
@@ -574,7 +569,7 @@ class MaskedGINE(nn.Module):
         class_weights=None,
     ):
         super().__init__()
-        # Use GineEncoder internally (so wrapper is used here too)
+        # Use GineEncoder internally
         self.encoder = GineEncoder(
             node_emb_dim=node_emb_dim,
             edge_emb_dim=edge_emb_dim,
@@ -595,7 +590,7 @@ class MaskedGINE(nn.Module):
             self.log_var_z = None
             self.log_var_pos = None
 
-        # expose class_weights same way as before
+        # class_weights 
         self.class_weights = getattr(self.encoder, "class_weights", None)
 
     def forward(
@@ -885,7 +880,7 @@ def train_and_evaluate(args: argparse.Namespace) -> None:
         except Exception as e:
             print(f"\nFailed to load best model from {best_model_path}: {e}")
 
-    # Final evaluation (same as your original)
+    # Final evaluation
     model.eval()
     preds_z_all, true_z_all = [], []
     pred_dists_all, true_dists_all = [], []