gat aug node pool class

Author: aditya0by0

chebai_graph/models/__init__.py

@@ -1,4 +1,5 @@
 from .augmented import (
+    GATAugNodePoolGraphPred,
     ResGatedAugNodePoolGraphPred,
     ResGatedAugOnlyPoolGraphPred,
     ResGatedFGNodeNoGraphNodeGraphPred,
@@ -24,4 +25,5 @@
     "ResGatedAugOnlyPoolGraphPred",
     "ResGatedGraphNodeOnlyPoolGraphPred",
     "ResGatedFGOnlyPoolGraphPred",
+    "GATAugNodePoolGraphPred",
 ]

chebai_graph/models/augmented.py

@@ -9,6 +9,7 @@
     GraphNodeOnlyPoolingNet,
     GraphNodePoolingNet,
 )
+from .gat import GATGraphPred
 from .resgated import ResGatedGraphPred
 
 
@@ -22,6 +23,16 @@ class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred):
     ...
 
 
+class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred):
+    """
+    Combines:
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings with molecule attributes.
+    - GATGraphPred: Graph attention network for final graph prediction.
+    """
+
+    ...
+
+
 class ResGatedGraphNodePoolGraphPred(GraphNodePoolingNet, ResGatedGraphPred):
     """
     Combines: