INNER CODE UNIT · Python

ncRNAClassificationWrapper

lbcb-sci/RiNALMo · train_ncrna_classification.py:26

class ncRNAClassificationWrapper(pl.LightningModule):
    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,

View source record →