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

Fixing empty prompts for text-generation when BOS exists. (#13859)

Browse Source 
* Fixing empty prompts for text-generation when BOS exists.

* Fixing odd case with Pegasus.

* Fixing Bert is Assertion Error.
This commit is contained in:
Nicolas Patry
2021-10-05 13:46:10 +02:00
committed by GitHub
parent a6ea244f99
commit 3a9c0f23b4
2 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/pipelines/text_generation.py
View File

@@ -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}

11
tests/test_pipelines_text_generation.py
View File

@@ -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("")