Commit
·
979833a
1
Parent(s):
e3a7bff
Update modeling_gcn.py
Browse files- 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 |
|