from transformers import PretrainedConfig


class RNATorsionBertConfig(PretrainedConfig):
    model_type = "rna_torsionbert"
    def __init__(self, k: int = 3, angles: str = "BACKBONE", hidden_size: int = 1024, **kwargs):
        """
        Initialise the model.
        :param k: the k-mer size.
        :param angles: types of angles to use. "BACKBONE" for eight torsional angles +
                                                            two pseudo-torsional angles
        :param hidden_size: size of the hidden layer after BERT hidden states.
        """
        self.k = k
        self.angles = angles
        self.hidden_size = hidden_size
        super().__init__(**kwargs)