INNER CODE UNIT · Python

ExpressionLevelWrapper

lbcb-sci/RiNALMo · train_expression_level.py:26

class ExpressionLevelWrapper(pl.LightningModule):
    def __init__(
        self,
        lm_config: str = "giga",
        head_embed_dim: int = 32,
        head_num_blocks: int = 6,
        lr: float = 1e-3,
    ):
        super().__init__()
        self.save_hyperparameters()

        self.lm = RiNALMo(model_config(lm_config))

        self.pred_head = RibosomeLoadingPredictionHead(
            c_in=self.lm.config['model']['transformer'].embed_dim,
            embed_dim=head_embed_dim,
            num_blocks=head_num_blocks
        )

View source record →