from transformers import PretrainedConfig

class SeqScreenConfig(PretrainedConfig):
  model_type = "seqscreen"
  def __init__(
    self,
    prot_dim: int = 2560, 
    mol_dim: int = 768, 
    proj_dim: int = 512, 
    dropout: float = 0.1,
    **kwargs):
    super().__init__(**kwargs)
    self.prot_dim = prot_dim
    self.mol_dim = mol_dim
    self.proj_dim = proj_dim
    self.dropout = dropout