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)