from transformers import PretrainedConfig

class ProtEnrichConfig(PretrainedConfig):
  model_type = "protenrich"

  def __init__(
    self,
    seq_dim : int = 1152, 
    struct_dim: int = 1024, 
    dyn_dim: int = 20, 
    embed_dim: int = 1024, 
    project_dim: int = 256, 
    alpha_max: float = 0.3,
    num_labels: int = 2,
    **kwargs):
    super().__init__(**kwargs)
    self.seq_dim = seq_dim
    self.struct_dim = struct_dim
    self.dyn_dim = dyn_dim
    self.embed_dim = embed_dim
    self.project_dim = project_dim
    self.alpha_max = alpha_max
    self.num_labels = num_labels