Enable TruncationStrategy override for pipelines (#9432)
* Enable TruncationStrategy override for pipelines * Update isort. * Fixing test * Fixing text_generation pipeline. * Using same DummyTok as other PR for easier merge later. * Some more import guards. * Remove bogus file. * Do not pass `generate_kwargs` to `_parse_and_tokenize`. @patrickvonplaten * Removed DummyTok. * Doc quality.
This commit is contained in:
@@ -14,15 +14,72 @@
|
||||
|
||||
import unittest
|
||||
|
||||
from transformers import pipeline
|
||||
from transformers import AutoTokenizer, is_torch_available, pipeline
|
||||
from transformers.testing_utils import require_torch, slow, torch_device
|
||||
from transformers.tokenization_utils import TruncationStrategy
|
||||
|
||||
from .test_pipelines_common import MonoInputPipelineCommonMixin
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
import torch
|
||||
|
||||
from transformers.models.bart import BartConfig, BartForConditionalGeneration
|
||||
|
||||
DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0
|
||||
|
||||
|
||||
class SimpleSummarizationPipelineTests(unittest.TestCase):
|
||||
@require_torch
|
||||
def test_input_too_long(self):
|
||||
torch.manual_seed(0)
|
||||
config = BartConfig(
|
||||
vocab_size=257,
|
||||
d_model=32,
|
||||
encoder_layers=1,
|
||||
decoder_layers=1,
|
||||
encoder_ffn_dim=32,
|
||||
decoder_ffn_dim=32,
|
||||
# So any text > 4 should raise an exception
|
||||
max_position_embeddings=4,
|
||||
encoder_attention_heads=1,
|
||||
decoder_attention_heads=1,
|
||||
max_length=4,
|
||||
min_length=1,
|
||||
)
|
||||
model = BartForConditionalGeneration(config)
|
||||
# Bias output towards L
|
||||
V, C = model.lm_head.weight.shape
|
||||
|
||||
bias = torch.zeros(V, requires_grad=True)
|
||||
bias[76] = 10
|
||||
|
||||
model.lm_head.bias = torch.nn.Parameter(bias)
|
||||
|
||||
# # Generated with:
|
||||
# import tempfile
|
||||
# from tokenizers import Tokenizer, models
|
||||
# from transformers import PreTrainedTokenizerFast
|
||||
# model_max_length = 4
|
||||
# vocab = [(chr(i), i) for i in range(256)]
|
||||
# tokenizer = Tokenizer(models.Unigram(vocab))
|
||||
# with tempfile.NamedTemporaryFile() as f:
|
||||
# tokenizer.save(f.name)
|
||||
# real_tokenizer = PreTrainedTokenizerFast(tokenizer_file=f.name, model_max_length=model_max_length)
|
||||
# real_tokenizer._tokenizer.save("tokenizer.json")
|
||||
# # + add missing config.json with albert as model_type
|
||||
tokenizer = AutoTokenizer.from_pretrained("Narsil/small_summarization_test")
|
||||
nlp = pipeline(task="summarization", model=model, tokenizer=tokenizer)
|
||||
|
||||
with self.assertLogs("transformers", level="WARNING"):
|
||||
with self.assertRaises(IndexError):
|
||||
_ = nlp("This is a test")
|
||||
|
||||
output = nlp("This is a test", truncation=TruncationStrategy.ONLY_FIRST)
|
||||
# 2 is default BOS from Bart.
|
||||
self.assertEqual(output, [{"summary_text": "\x02 L L L"}])
|
||||
|
||||
|
||||
class SummarizationPipelineTests(MonoInputPipelineCommonMixin, unittest.TestCase):
|
||||
pipeline_task = "summarization"
|
||||
pipeline_running_kwargs = {"num_beams": 2, "min_length": 2, "max_length": 5}
|
||||
|
||||
Reference in New Issue
Block a user