| from datasets.load import load_dataset | |
| import logging | |
| import sacrebleu | |
| import pandas as pd | |
| from simpletransformers.t5 import T5Model, T5Args | |
| raw_datasets = load_dataset('iwslt2017', 'iwslt2017-zh-en') | |
| logging.basicConfig(level=logging.INFO) | |
| transformers_logger = logging.getLogger("transformers") | |
| transformers_logger.setLevel(logging.WARNING) | |
| model_args = T5Args() | |
| model_args.max_length = 512 | |
| model_args.length_penalty = 1 | |
| model_args.num_beams = 10 | |
| model = T5Model("mt5", "outputs", args=model_args) | |
| en_zh_test = pd.DataFrame(raw_datasets['test']['translation']) | |
| zh_truth = en_zh_test['zh'].tolist() | |
| en_input = en_zh_test['en'].tolist() | |
| zh_preds = model.predict(en_input) | |
| en_zh_bleu = sacrebleu.corpus_bleu(zh_preds, zh_truth) | |
| print("----------------------------------------------") | |
| print("English to Chinese: ", en_zh_bleu.score) | |
| en_preds = model.predict(zh_truth) | |
| zh_en_bleu = sacrebleu.corpus_bleu(en_preds, en_input) | |
| print("----------------------------------------------") | |
| print("Chinese to English: ", zh_en_bleu.score) | |