Huhujingjing
/

custom-gcn

Graph Machine Learning

feature-extraction

Model card Files Files and versions Community

custom-gcn / configuration_gcn.py

Huhujingjing's picture

Upload model

2444fad 12 months ago

raw history blame contribute delete

1.21 kB

	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")