INNER CODE UNIT · Python

main

Adibvafa/CodonTransformer · finetune.py:126

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
    train_data = IterableJSONData(args.dataset_dir, dist_env="slurm")
    data_loader = DataLoader(
        dataset=train_data,
        collate_fn=MaskedTokenizerCollator(tokenizer),
        batch_size=args.batch_size,
        num_workers=0 if args.debug else args.num_workers,
        persistent_workers=False if args.debug else True,

View source record →