From 956c4c4eb4e175ddf7be0a9ed08f5f4e5a84c5e9 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Tue, 19 May 2020 19:45:31 -0400 Subject: [PATCH] [gpu slow tests] fix mbart-large-enro gpu tests (#4472) --- tests/test_modeling_bart.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tests/test_modeling_bart.py b/tests/test_modeling_bart.py index 0724e18efd..718d41ca30 100644 --- a/tests/test_modeling_bart.py +++ b/tests/test_modeling_bart.py @@ -231,7 +231,7 @@ class BartTranslationTests(unittest.TestCase): """Only load the model if needed.""" if self._model is None: model = BartForConditionalGeneration.from_pretrained("mbart-large-en-ro") - self._model = model + self._model = model.to(torch_device) return self._model @slow @@ -257,10 +257,7 @@ class BartTranslationTests(unittest.TestCase): ) } translated_tokens = model.generate(input_ids=inputs["input_ids"].to(torch_device), num_beams=5,) - decoded = [ - self.tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) - for g in translated_tokens - ] + decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True) self.assertEqual(expected_translation_romanian, decoded[0]) def test_mbart_enro_config(self): @@ -576,11 +573,13 @@ class BartModelIntegrationTests(unittest.TestCase): PGE_ARTICLE = """ PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow.""" EXPECTED_SUMMARY = "California's largest power company has begun shutting off power to tens of thousands of homes and businesses in the state." - dct = tok.batch_encode_plus([PGE_ARTICLE], max_length=1024, pad_to_max_length=True, return_tensors="pt",) + dct = tok.batch_encode_plus([PGE_ARTICLE], max_length=1024, pad_to_max_length=True, return_tensors="pt",).to( + torch_device + ) hypotheses_batch = model.generate( - input_ids=dct["input_ids"].to(torch_device), - attention_mask=dct["attention_mask"].to(torch_device), + input_ids=dct["input_ids"], + attention_mask=dct["attention_mask"], num_beams=2, max_length=62, min_length=11, @@ -590,9 +589,7 @@ class BartModelIntegrationTests(unittest.TestCase): decoder_start_token_id=model.config.eos_token_id, ) - decoded = [ - tok.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in hypotheses_batch - ] + decoded = tok.batch_decode(hypotheses_batch, skip_special_tokens=True,) self.assertEqual(EXPECTED_SUMMARY, decoded[0]) def test_xsum_config_generation_params(self):