Unverified Commit 3a9c0f23 authored by Nicolas Patry's avatar Nicolas Patry Committed by GitHub
Browse files

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.
parent a6ea244f
...@@ -158,6 +158,9 @@ class TextGenerationPipeline(Pipeline): ...@@ -158,6 +158,9 @@ class TextGenerationPipeline(Pipeline):
def _forward(self, model_inputs, **generate_kwargs): def _forward(self, model_inputs, **generate_kwargs):
input_ids = model_inputs["input_ids"] 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") prompt_text = model_inputs.pop("prompt_text")
generated_sequence = self.model.generate(input_ids=input_ids, **generate_kwargs) # BS x SL 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} return {"generated_sequence": generated_sequence, "input_ids": input_ids, "prompt_text": prompt_text}
......
...@@ -106,3 +106,14 @@ class TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseM ...@@ -106,3 +106,14 @@ class TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseM
outputs = text_generator("This is a test", return_full_text=True) outputs = text_generator("This is a test", return_full_text=True)
self.assertEqual(outputs, [{"generated_text": ANY(str)}]) self.assertEqual(outputs, [{"generated_text": ANY(str)}])
self.assertTrue(outputs[0]["generated_text"].startswith("This is a test")) 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("")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment