INNER CODE UNIT · Python

main

Adibvafa/CodonTransformer · pretrain.py:128

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

    # Load the tokenizer and model
    tokenizer = PreTrainedTokenizerFast(
        tokenizer_file=args.tokenizer_path,
        bos_token="[CLS]",
        eos_token="[SEP]",
        unk_token="[UNK]",
        sep_token="[SEP]",
        pad_token="[PAD]",
        cls_token="[CLS]",
        mask_token="[MASK]",
    )
    config = BigBirdConfig(
        vocab_size=len(tokenizer),

View source record →