Update modeling_transmxm.py

modeling_transmxm.py

@@ -23,7 +23,7 @@ from torch_scatter import scatter
 from torch_geometric.nn import global_add_pool, radius
 from torch_sparse import SparseTensor
 
-from transmxm_model.configuration_transmxm import TransmxmConfig
+# from transmxm_model.configuration_transmxm import TransmxmConfig
 
 from tqdm import tqdm
 import numpy as np
@@ -1171,35 +1171,35 @@ class Local_MP(torch.nn.Module):
         return h, y
 
 
-# class MXMConfig(PretrainedConfig):
-#     model_type = "gcn"
-#
-#     def __init__(
-#         self,
-#         dim: int=128,
-#         n_layer: int=6,
-#         cutoff: float=5.0,
-#         num_spherical: int=7,
-#         num_radial: int=6,
-#         envelope_exponent: int=5,
-#
-#         smiles: List[str] = None,
-#         processor_class: str = "SmilesProcessor",
-#         **kwargs,
-#     ):
-#
-#         self.dim = dim                                    # the dimension of input feature
-#         self.n_layer = n_layer                            # the number of GCN layers
-#         self.cutoff = cutoff                              # the cutoff distance for neighbor searching
-#         self.num_spherical = num_spherical                # the number of spherical harmonics
-#         self.num_radial = num_radial                         # the number of radial basis
-#         self.envelope_exponent = envelope_exponent             # the envelope exponent
-#
-#         self.smiles = smiles                      # process smiles
-#         self.processor_class = processor_class
-#
-#
-#         super().__init__(**kwargs)
+class TransmxmConfig(PretrainedConfig):
+    model_type = "transmxm"
+
+    def __init__(
+        self,
+        dim: int=128,
+        n_layer: int=6,
+        cutoff: float=5.0,
+        num_spherical: int=7,
+        num_radial: int=6,
+        envelope_exponent: int=5,
+
+        smiles: List[str] = None,
+        processor_class: str = "SmilesProcessor",
+        **kwargs,
+    ):
+
+        self.dim = dim                                    # the dimension of input feature
+        self.n_layer = n_layer                            # the number of GCN layers
+        self.cutoff = cutoff                              # the cutoff distance for neighbor searching
+        self.num_spherical = num_spherical                # the number of spherical harmonics
+        self.num_radial = num_radial                         # the number of radial basis
+        self.envelope_exponent = envelope_exponent             # the envelope exponent
+
+        self.smiles = smiles                      # process smiles
+        self.processor_class = processor_class
+
+
+        super().__init__(**kwargs)