Upload model

modeling_gcn.py

@@ -1,47 +1,24 @@
-from torch_geometric.nn import GCNConv
+import os
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_scatter import scatter
-from transformers import PreTrainedModel
-from gcn_model.configuration_gcn import GCNConfig
-import torch
-from rdkit import Chem
-from rdkit.Chem import AllChem
-import torch
-from torch_geometric.data import Data
-import os
-from transformers import PretrainedConfig
-from typing import List
-from torch_geometric.loader import DataLoader
 from tqdm import tqdm
 import pandas as pd
-from transformers import AutoModel
-class GCNConfig(PretrainedConfig):
-    model_type = "gcn"
+from typing import List
 
-    def __init__(
-        self,
-        input_feature: int=64,
-        emb_input: int=20,
-        hidden_size: int=64,
-        n_layers: int=6,
-        num_classes: int=1,
+from rdkit import Chem
+from rdkit.Chem import AllChem
 
-        smiles: List[str] = None,
-        processor_class: str = "SmilesProcessor",
-        **kwargs,
-    ):
+from transformers import PretrainedConfig
+from transformers import PreTrainedModel
+from transformers import AutoModel
 
-        self.input_feature = input_feature        # the dimension of input feature
-        self.emb_input = emb_input                # the embedding dimension of input feature
-        self.hidden_size = hidden_size            # the hidden size of GCN
-        self.n_layers = n_layers                  # the number of GCN layers
-        self.num_classes = num_classes            # the number of output classes
+from torch_geometric.nn import GCNConv
+from torch_geometric.data import Data
+from torch_geometric.loader import DataLoader
+from torch_scatter import scatter
 
-        self.smiles = smiles                      # process smiles
-        self.processor_class = processor_class
 
-        super().__init__(**kwargs)
 class SmilesDataset(torch.utils.data.Dataset):
     def __init__(self, smiles):
         self.smiles_list = smiles
@@ -176,6 +153,35 @@ class GCNNet(torch.nn.Module):
 
         return x.squeeze(-1)
 
+
+class GCNConfig(PretrainedConfig):
+    model_type = "gcn"
+
+    def __init__(
+        self,
+        input_feature: int=64,
+        emb_input: int=20,
+        hidden_size: int=64,
+        n_layers: int=6,
+        num_classes: int=1,
+
+        smiles: List[str] = None,
+        processor_class: str = "SmilesProcessor",
+        **kwargs,
+    ):
+
+        self.input_feature = input_feature        # the dimension of input feature
+        self.emb_input = emb_input                # the embedding dimension of input feature
+        self.hidden_size = hidden_size            # the hidden size of GCN
+        self.n_layers = n_layers                  # the number of GCN layers
+        self.num_classes = num_classes            # the number of output classes
+
+        self.smiles = smiles                      # process smiles
+        self.processor_class = processor_class
+
+        super().__init__(**kwargs)
+
+
 class GCNModel(PreTrainedModel):
     config_class = GCNConfig