INNER CODE UNIT · Python
__init__
lbcb-sci/RiNALMo · train_ncrna_classification.py:27
def __init__(
self,
lm_config: str = "giga",
head_embed_dim: int = 256,
n_classes: int = 88,
lr: float = 8e-6,
weight_decay: float = 0.01,
) -> None:
super().__init__()
self.save_hyperparameters()
self.rinalmo = RiNALMo(model_config(lm_config))
self.pred_head = ncRNAClassificationHead(
c_in=self.rinalmo.config['model']['transformer'].embed_dim,
embed_dim=head_embed_dim,
n_classes=n_classes,
)