model.py

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

def create_mt5_small():
  """
  Initializes model and tokenizer.
  """
  checkpoint = 'google/mt5-small'
  tokenizer = AutoTokenizer.from_pretrained(checkpoint, return_tensors='pt')
  model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)

  return model, tokenizer