File size: 1,798 Bytes
5cedadf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
from datasets.load import load_dataset
import pandas as pd
import logging
from simpletransformers.t5 import T5Args, T5Model
logging.basicConfig(level=logging.INFO)
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)
raw_datasets = load_dataset('iwslt2017', 'iwslt2017-zh-en')
train_df = pd.DataFrame(raw_datasets['train']['translation'])
train_df.columns = ['input_text', 'target_text']
reverse_df = train_df.copy()
reverse_df.columns = ['target_text', 'input_text']
train_df['prefix'] = 'translate english to chinese'
reverse_df['prefix'] = 'translate chinese to english'
train_df = pd.concat([train_df, reverse_df])
eval_df = pd.DataFrame(raw_datasets['validation']['translation'])
eval_df.columns = ['input_text', 'target_text']
reverse_df = eval_df.copy()
reverse_df.columns = ['target_text', 'input_text']
eval_df['prefix'] = 'translate english to chinese'
reverse_df['prefix'] = 'translate chinese to english'
eval_df = pd.concat([eval_df, reverse_df])
model_args = T5Args()
model_args.max_seq_length = 96
model_args.train_batch_size = 20
model_args.eval_batch_size = 20
model_args.num_train_epochs = 4
model_args.evaluate_during_training = True
model_args.evaluate_during_training_steps = 5000
model_args.use_multiprocessing = False
model_args.fp16 = False
model_args.save_steps = -1
model_args.save_model_every_epoch = True
model_args.save_eval_checkpoints = False
model_args.no_cache = True
model_args.reprocess_input_data = True
model_args.overwrite_output_dir = False
model_args.preprocess_inputs = False
model_args.num_return_sequences = 1
model_args.wandb_project = "MT5 English-Chinese Translation"
model = T5Model("mt5", "outputs", args=model_args)
model.train_model(train_df, eval_data=eval_df, output_dir='mt5_more_epochs')
|