[Examples] Clean summarization and translation example testing files for T5 and Bart (#3514)

* fix conflicts

* add model size argument to summarization

* correct wrong import

* fix isort

* correct imports

* other isort make style

* make style
This commit is contained in:
Patrick von Platen authored and GitHub committed 2020-03-31 17:54:13 +02:00
1 parent 0373b60c4c
commit ae6834e028
6 files changed
+36 -22

No files matched your search

+10 -4
View File
@@ -14,13 +14,13 @@ def chunks(lst, n):
yield lst[i : i + n]
def generate_translations(lns, output_file_path, batch_size, device):
def generate_translations(lns, output_file_path, model_size, batch_size, device):
output_file = Path(output_file_path).open("w")
model = T5ForConditionalGeneration.from_pretrained("t5-base")
model = T5ForConditionalGeneration.from_pretrained(model_size)
model.to(device)
tokenizer = T5Tokenizer.from_pretrained("t5-base")
tokenizer = T5Tokenizer.from_pretrained(model_size)
# update config with summarization specific params
task_specific_params = model.config.task_specific_params
@@ -52,6 +52,12 @@ def calculate_bleu_score(output_lns, refs_lns, score_path):
def run_generate():
parser = argparse.ArgumentParser()
parser.add_argument(
"model_size",
type=str,
help="T5 model size, either 't5-small', 't5-base', 't5-large', 't5-3b', 't5-11b'. Defaults to 't5-base'.",
default="t5-base",
)
parser.add_argument(
"input_path", type=str, help="like wmt/newstest2013.en",
)
@@ -78,7 +84,7 @@ def run_generate():
input_lns = [x.strip().replace(dash_pattern[0], dash_pattern[1]) for x in open(args.input_path).readlines()]
generate_translations(input_lns, args.output_path, args.batch_size, args.device)
generate_translations(input_lns, args.output_path, args.model_size, args.batch_size, args.device)
output_lns = [x.strip() for x in open(args.output_path).readlines()]
refs_lns = [x.strip().replace(dash_pattern[0], dash_pattern[1]) for x in open(args.reference_path).readlines()]
+11 -4
View File
@@ -1,5 +1,4 @@
import logging
import os
import sys
import tempfile
import unittest
@@ -33,11 +32,19 @@ class TestT5Examples(unittest.TestCase):
with tmp_target.open("w") as f:
f.write("\n".join(translation))
testargs = ["evaluate_wmt.py", str(tmp_source), output_file_name, str(tmp_target), score_file_name]
output_file_name = Path(tempfile.gettempdir()) / "utest_output_trans.hypo"
score_file_name = Path(tempfile.gettempdir()) / "utest_score.hypo"
testargs = [
"evaluate_wmt.py",
"t5-small",
str(tmp_source),
str(output_file_name),
str(tmp_target),
str(score_file_name),
]
with patch.object(sys, "argv", testargs):
run_generate()
self.assertTrue(Path(output_file_name).exists())
self.assertTrue(Path(score_file_name).exists())
os.remove(Path(output_file_name))
os.remove(Path(score_file_name))