SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Add option to choose T5 model size. (#3480)

Browse Source 
T5-small in test


isort
This commit is contained in:
Lysandre Debut
2020-03-27 10:56:59 -04:00
committed by GitHub
parent e2c05f06ef
commit ff80b73157
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
examples/summarization/t5/evaluate_cnn.py
View File

@@ -14,13 +14,13 @@ def chunks(lst, n):
yield lst[i : i + n] yield lst[i : i + n]
def generate_summaries(lns, output_file_path, batch_size, device): def generate_summaries(lns, output_file_path, model_size, batch_size, device):
output_file = Path(output_file_path).open("w") output_file = Path(output_file_path).open("w")
model = T5ForConditionalGeneration.from_pretrained("t5-large") model = T5ForConditionalGeneration.from_pretrained(model_size)
model.to(device) model.to(device)
tokenizer = T5Tokenizer.from_pretrained("t5-large") tokenizer = T5Tokenizer.from_pretrained(model_size)
# update config with summarization specific params # update config with summarization specific params
task_specific_params = model.config.task_specific_params task_specific_params = model.config.task_specific_params
@@ -61,6 +61,12 @@ def calculate_rouge(output_lns, reference_lns, score_path):
def run_generate(): def run_generate():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument(
"model_size",
type=str,
help="T5 model size, either 't5-small', 't5-base' or 't5-large'. Defaults to base.",
default="t5-base",
)
parser.add_argument( parser.add_argument(
"input_path", type=str, help="like cnn_dm/test_articles_input.txt", "input_path", type=str, help="like cnn_dm/test_articles_input.txt",
) )
@@ -83,7 +89,7 @@ def run_generate():
source_lns = [x.rstrip() for x in open(args.input_path).readlines()] source_lns = [x.rstrip() for x in open(args.input_path).readlines()]
generate_summaries(source_lns, args.output_path, args.batch_size, args.device) generate_summaries(source_lns, args.output_path, args.model_size, args.batch_size, args.device)
output_lns = [x.rstrip() for x in open(args.output_path).readlines()] output_lns = [x.rstrip() for x in open(args.output_path).readlines()]
reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()] reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()]

2
examples/summarization/t5/test_t5_examples.py
View File

@@ -22,7 +22,7 @@ class TestT5Examples(unittest.TestCase):
tmp = Path(tempfile.gettempdir()) / "utest_generations.hypo" tmp = Path(tempfile.gettempdir()) / "utest_generations.hypo"
with tmp.open("w") as f: with tmp.open("w") as f:
f.write("\n".join(articles)) f.write("\n".join(articles))
testargs = ["evaluate_cnn.py", str(tmp), "output.txt", str(tmp), "score.txt"] testargs = ["evaluate_cnn.py", "t5-small", str(tmp), "output.txt", str(tmp), "score.txt"]
with patch.object(sys, "argv", testargs): with patch.object(sys, "argv", testargs):
run_generate() run_generate()
self.assertTrue(Path("output.txt").exists()) self.assertTrue(Path("output.txt").exists())