INNER CODE UNIT · Python

on_save_checkpoint

Adibvafa/CodonTransformer · finetune.py:119

    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
        model = trainer.model.model
        torch.save(
            model.state_dict(), os.path.join(self.dirpath, self.checkpoint_filename)
        )


def main(args):
    """Finetune the CodonTransformer model."""
    pl.seed_everything(args.seed)
    torch.set_float32_matmul_precision("medium")

    # Load the tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained("adibvafa/CodonTransformer")
    model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer-base")
    harnessed_model = plTrainHarness(model, args.learning_rate, args.warmup_fraction)

    # Load the training data

View source record →