Huhujingjing commited on
Commit
979833a
·
1 Parent(s): e3a7bff

Update modeling_gcn.py

Browse files
Files changed (1) hide show
  1. modeling_gcn.py +30 -0
modeling_gcn.py CHANGED
@@ -145,6 +145,36 @@ class GCNNet(torch.nn.Module):
145
 
146
  return x.squeeze(-1)
147
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
148
  class GCNModel(PreTrainedModel):
149
  config_class = GCNConfig
150
 
 
145
 
146
  return x.squeeze(-1)
147
 
148
+
149
+ from transformers import PretrainedConfig
150
+ from typing import List
151
+ class GCNConfig(PretrainedConfig):
152
+ model_type = "gcn"
153
+
154
+ def __init__(
155
+ self,
156
+ input_feature: int=64,
157
+ emb_input: int=20,
158
+ hidden_size: int=64,
159
+ n_layers: int=6,
160
+ num_classes: int=1,
161
+ smiles: List[str] = None,
162
+ processor_class: str = "SmilesProcessor",
163
+ **kwargs,
164
+ ):
165
+
166
+ self.input_feature = input_feature # the dimension of input feature
167
+ self.emb_input = emb_input # the embedding dimension of input feature
168
+ self.hidden_size = hidden_size # the hidden size of GCN
169
+ self.n_layers = n_layers # the number of GCN layers
170
+ self.num_classes = num_classes # the number of output classes
171
+
172
+ self.smiles = smiles # process smiles
173
+ self.processor_class = processor_class
174
+
175
+ super().__init__(**kwargs)
176
+
177
+
178
  class GCNModel(PreTrainedModel):
179
  config_class = GCNConfig
180