from transformers import PretrainedConfig | |
class RNATorsionBertConfig(PretrainedConfig): | |
model_type = "rna_torsionbert" | |
def __init__(self, k: int = 3, num_classes: int = 18, hidden_size: int = 1024, **kwargs): | |
""" | |
Initialise the model. | |
:param k: the k-mer size. | |
:param num_classes: the number of labels. | |
:param hidden_size: size of the hidden layer after BERT hidden states. | |
""" | |
self.k = k | |
self.num_classes = num_classes | |
self.hidden_size = hidden_size | |
super().__init__(**kwargs) |