File size: 1,205 Bytes
2444fad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
from transformers import PretrainedConfig
from typing import List
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)
if __name__ == "__main__":
gcn_config = GCNConfig(input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1, smiles=["C", "CC", "CCC"], processor_class="SmilesProcessor")
gcn_config.save_pretrained("custom-gcn") |