main.py

import deepspeed
import sentencepiece as spm

from arguments import get_args
from config import ModelArgs, TrainArgs
from data import TextFileDataset, create_dataloaders
from model import Transformer
from train import train


def run_main():
      args = get_args()
      tokenizer = spm.SentencePieceProcessor()
      tokenizer.load(args['tokenizer_model_path'])
      dataset = TextFileDataset(args['filepath'], tokenizer, args['max_seq_length'], pad_token_id=0)
      model_args = ModelArgs()
      model_args.vocab_size = tokenizer.vocab_size()
      if args['deepspeed']:
            deepspeed.init_distributed()
      train_dataloader, eval_dataloader = create_dataloaders(dataset, args['train_size'], args['batch_size'])

      model = Transformer(model_args)
      train(model=model,
            train_config=TrainArgs(),
            train_dataloader=train_dataloader,
            eval_dataloader=eval_dataloader,
            args=args)


if __name__ == '__main__':
      run_main()