"examples/vscode:/vscode.git/clone" did not exist on "b8344a274fe13b390fa60c74b76117f5ea8144cb"
Unverified Commit e13a91fe authored by Nicolas Patry's avatar Nicolas Patry Committed by GitHub
Browse files

Fixing return type tensor with `num_return_sequences>1`. (#16828)

* Fixing return type tensor with `num_return_sequences>1`.

* Nit.
parent ff06b177
...@@ -168,7 +168,7 @@ class Text2TextGenerationPipeline(Pipeline): ...@@ -168,7 +168,7 @@ class Text2TextGenerationPipeline(Pipeline):
records = [] records = []
for output_ids in model_outputs["output_ids"][0]: for output_ids in model_outputs["output_ids"][0]:
if return_type == ReturnType.TENSORS: if return_type == ReturnType.TENSORS:
record = {f"{self.return_name}_token_ids": model_outputs} record = {f"{self.return_name}_token_ids": output_ids}
elif return_type == ReturnType.TEXT: elif return_type == ReturnType.TEXT:
record = { record = {
f"{self.return_name}_text": self.tokenizer.decode( f"{self.return_name}_text": self.tokenizer.decode(
......
...@@ -226,7 +226,7 @@ class TextGenerationPipeline(Pipeline): ...@@ -226,7 +226,7 @@ class TextGenerationPipeline(Pipeline):
records = [] records = []
for sequence in generated_sequence: for sequence in generated_sequence:
if return_type == ReturnType.TENSORS: if return_type == ReturnType.TENSORS:
record = {"generated_token_ids": generated_sequence} record = {"generated_token_ids": sequence}
elif return_type in {ReturnType.NEW_TEXT, ReturnType.FULL_TEXT}: elif return_type in {ReturnType.NEW_TEXT, ReturnType.FULL_TEXT}:
# Decode text # Decode text
text = self.tokenizer.decode( text = self.tokenizer.decode(
......
...@@ -21,10 +21,15 @@ from transformers import ( ...@@ -21,10 +21,15 @@ from transformers import (
pipeline, pipeline,
) )
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available
from .test_pipelines_common import ANY, PipelineTestCaseMeta from .test_pipelines_common import ANY, PipelineTestCaseMeta
if is_torch_available():
import torch
@is_pipeline_test @is_pipeline_test
class Text2TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseMeta): class Text2TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseMeta):
model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
...@@ -83,6 +88,37 @@ class Text2TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTest ...@@ -83,6 +88,37 @@ class Text2TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTest
] ]
self.assertEqual(outputs, target_outputs) self.assertEqual(outputs, target_outputs)
outputs = generator("This is a test", do_sample=True, num_return_sequences=2, return_tensors=True)
self.assertEqual(
outputs,
[
{"generated_token_ids": ANY(torch.Tensor)},
{"generated_token_ids": ANY(torch.Tensor)},
],
)
generator.tokenizer.pad_token_id = generator.model.config.eos_token_id
generator.tokenizer.pad_token = "<pad>"
outputs = generator(
["This is a test", "This is a second test"],
do_sample=True,
num_return_sequences=2,
batch_size=2,
return_tensors=True,
)
self.assertEqual(
outputs,
[
[
{"generated_token_ids": ANY(torch.Tensor)},
{"generated_token_ids": ANY(torch.Tensor)},
],
[
{"generated_token_ids": ANY(torch.Tensor)},
{"generated_token_ids": ANY(torch.Tensor)},
],
],
)
@require_tf @require_tf
def test_small_model_tf(self): def test_small_model_tf(self):
generator = pipeline("text2text-generation", model="patrickvonplaten/t5-tiny-random", framework="tf") generator = pipeline("text2text-generation", model="patrickvonplaten/t5-tiny-random", framework="tf")
......
...@@ -56,6 +56,37 @@ class TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseM ...@@ -56,6 +56,37 @@ class TextGenerationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseM
], ],
) )
outputs = text_generator("This is a test", do_sample=True, num_return_sequences=2, return_tensors=True)
self.assertEqual(
outputs,
[
{"generated_token_ids": ANY(list)},
{"generated_token_ids": ANY(list)},
],
)
text_generator.tokenizer.pad_token_id = text_generator.model.config.eos_token_id
text_generator.tokenizer.pad_token = "<pad>"
outputs = text_generator(
["This is a test", "This is a second test"],
do_sample=True,
num_return_sequences=2,
batch_size=2,
return_tensors=True,
)
self.assertEqual(
outputs,
[
[
{"generated_token_ids": ANY(list)},
{"generated_token_ids": ANY(list)},
],
[
{"generated_token_ids": ANY(list)},
{"generated_token_ids": ANY(list)},
],
],
)
@require_tf @require_tf
def test_small_model_tf(self): def test_small_model_tf(self):
text_generator = pipeline(task="text-generation", model="sshleifer/tiny-ctrl", framework="tf") text_generator = pipeline(task="text-generation", model="sshleifer/tiny-ctrl", framework="tf")
......
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