From 3a9c0f23b47311f373d890bca1b13aa99fb55a34 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 5 Oct 2021 13:46:10 +0200 Subject: [PATCH] Fixing empty prompts for text-generation when BOS exists. (#13859) * Fixing empty prompts for text-generation when BOS exists. * Fixing odd case with Pegasus. * Fixing Bert is Assertion Error. --- src/transformers/pipelines/text_generation.py | 3 +++ tests/test_pipelines_text_generation.py | 11 +++++++++++ 2 files changed, 14 insertions(+) diff --git a/src/transformers/pipelines/text_generation.py b/src/transformers/pipelines/text_generation.py index 53c626297d..03d9621b4c 100644 --- a/src/transformers/pipelines/text_generation.py +++ b/src/transformers/pipelines/text_generation.py @@ -158,6 +158,9 @@ class TextGenerationPipeline(Pipeline): def _forward(self, model_inputs, **generate_kwargs): input_ids = model_inputs["input_ids"] + # Allow empty prompts + if input_ids.shape[1] == 0: + input_ids = None prompt_text = model_inputs.pop("prompt_text") generated_sequence = self.model.generate(input_ids=input_ids, **generate_kwargs) # BS x SL return {"generated_sequence": generated_sequence, "input_ids": input_ids, "prompt_text": prompt_text} diff --git a/tests/test_pipelines_text_generation.py b/tests/test_pipelines_text_generation.py index 22bc8bf42c..3618a2be73 100644 --- a/tests/test_pipelines_text_generation.py +++ b/tests/test_pipelines_text_generation.py @@ -106,3 +106,14 @@ class TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseM outputs = text_generator("This is a test", return_full_text=True) self.assertEqual(outputs, [{"generated_text": ANY(str)}]) self.assertTrue(outputs[0]["generated_text"].startswith("This is a test")) + + # Empty prompt is slighly special + # it requires BOS token to exist. + # Special case for Pegasus which will always append EOS so will + # work even without BOS. + if text_generator.tokenizer.bos_token_id is not None or "Pegasus" in tokenizer.__class__.__name__: + outputs = text_generator("") + self.assertEqual(outputs, [{"generated_text": ANY(str)}]) + else: + with self.assertRaises((ValueError, AssertionError)): + outputs = text_generator("")