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.

src/transformers/pipelines/base.py

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from ..file_utils import add_end_docstrings, is_tf_available, is_torch_available
 from ..modelcard import ModelCard
-from ..tokenization_utils import PreTrainedTokenizer
+from ..tokenization_utils import PreTrainedTokenizer, TruncationStrategy
 from ..utils import logging
 
 
@@ -577,7 +577,9 @@ class Pipeline(_ScikitCompat):
                 f"The model '{self.model.__class__.__name__}' is not supported for {self.task}. Supported models are {supported_models}",
             )
 
-    def _parse_and_tokenize(self, inputs, padding=True, add_special_tokens=True, **kwargs):
+    def _parse_and_tokenize(
+        self, inputs, padding=True, add_special_tokens=True, truncation=TruncationStrategy.DO_NOT_TRUNCATE, **kwargs
+    ):
         """
         Parse arguments and tokenize
         """
@@ -587,6 +589,7 @@ class Pipeline(_ScikitCompat):
             add_special_tokens=add_special_tokens,
             return_tensors=self.framework,
             padding=padding,
+            truncation=truncation,
         )
 
         return inputs

src/transformers/pipelines/conversational.py

@@ -2,6 +2,7 @@ import uuid
 from typing import List, Optional, Union
 
 from ..file_utils import add_end_docstrings, is_tf_available, is_torch_available
+from ..tokenization_utils import TruncationStrategy
 from ..utils import logging
 from .base import PIPELINE_INIT_ARGS, Pipeline
 
@@ -317,12 +318,14 @@ class ConversationalPipeline(Pipeline):
             else:
                 return output
 
-    def _parse_and_tokenize(self, inputs, **kwargs):
+    def _parse_and_tokenize(
+        self, inputs, add_special_tokens=False, padding=False, truncation=TruncationStrategy.DO_NOT_TRUNCATE, **kwargs
+    ):
         """
         Parse arguments and tokenize, adding an EOS token at the end of the user input
         """
         # Parse arguments
-        inputs = self.tokenizer(inputs, add_special_tokens=False, padding=False).get("input_ids", [])
+        inputs = self.tokenizer(inputs, add_special_tokens=add_special_tokens, padding=padding).get("input_ids", [])
         for input in inputs:
             input.append(self.tokenizer.eos_token_id)
         return inputs

src/transformers/pipelines/text2text_generation.py

@@ -1,4 +1,5 @@
 from ..file_utils import add_end_docstrings, is_tf_available, is_torch_available
+from ..tokenization_utils import TruncationStrategy
 from ..utils import logging
 from .base import PIPELINE_INIT_ARGS, Pipeline
 
@@ -50,7 +51,13 @@ class Text2TextGenerationPipeline(Pipeline):
         return True
 
     def __call__(
-        self, *args, return_tensors=False, return_text=True, clean_up_tokenization_spaces=False, **generate_kwargs
+        self,
+        *args,
+        return_tensors=False,
+        return_text=True,
+        clean_up_tokenization_spaces=False,
+        truncation=TruncationStrategy.DO_NOT_TRUNCATE,
+        **generate_kwargs
     ):
         r"""
         Generate the output text(s) using text(s) given as inputs.
@@ -64,6 +71,10 @@ class Text2TextGenerationPipeline(Pipeline):
                 Whether or not to include the decoded texts in the outputs.
             clean_up_tokenization_spaces (:obj:`bool`, `optional`, defaults to :obj:`False`):
                 Whether or not to clean up the potential extra spaces in the text output.
+            truncation (:obj:`TruncationStrategy`, `optional`, defaults to :obj:`TruncationStrategy.DO_NOT_TRUNCATE`):
+                The truncation strategy for the tokenization within the pipeline.
+                :obj:`TruncationStrategy.DO_NOT_TRUNCATE` (default) will never truncate, but it is sometimes desirable
+                to truncate the input to fit the model's max_length instead of throwing an error down the line.
             generate_kwargs:
                 Additional keyword arguments to pass along to the generate method of the model (see the generate method
                 corresponding to your framework `here <./model.html#generative-models>`__).
@@ -96,7 +107,7 @@ class Text2TextGenerationPipeline(Pipeline):
             )
 
         with self.device_placement():
-            inputs = self._parse_and_tokenize(*args, padding=padding, **generate_kwargs)
+            inputs = self._parse_and_tokenize(*args, padding=padding, truncation=truncation)
 
             if self.framework == "pt":
                 inputs = self.ensure_tensor_on_device(**inputs)
@@ -108,9 +119,6 @@ class Text2TextGenerationPipeline(Pipeline):
             max_length = generate_kwargs.get("max_length", self.model.config.max_length)
             self.check_inputs(input_length, min_length, max_length)
 
-            # truncation should be used by _parse_and_tokenize
-            generate_kwargs.pop("truncation", None)
-
             generations = self.model.generate(
                 inputs["input_ids"],
                 attention_mask=inputs["attention_mask"],

src/transformers/pipelines/text_generation.py

@@ -50,25 +50,15 @@ class TextGenerationPipeline(Pipeline):
         self.check_model_type(self.ALLOWED_MODELS)
 
     # overriding _parse_and_tokenize to allow for unusual language-modeling tokenizer arguments
-
-    def _parse_and_tokenize(self, inputs, padding=True, add_special_tokens=True, **kwargs):
+    def _parse_and_tokenize(self, *args, **kwargs):
         """
         Parse arguments and tokenize
         """
         # Parse arguments
         if self.model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
-            tokenizer_kwargs = {"add_space_before_punct_symbol": True}
-        else:
-            tokenizer_kwargs = {}
-        inputs = self.tokenizer(
-            inputs,
-            add_special_tokens=add_special_tokens,
-            return_tensors=self.framework,
-            padding=padding,
-            **tokenizer_kwargs,
-        )
+            kwargs.update({"add_space_before_punct_symbol": True})
 
-        return inputs
+        return super()._parse_and_tokenize(*args, **kwargs)
 
     def __call__(
         self,

src/transformers/pipelines/zero_shot_classification.py

@@ -3,6 +3,7 @@ from typing import List, Union
 import numpy as np
 
 from ..file_utils import add_end_docstrings
+from ..tokenization_utils import TruncationStrategy
 from ..utils import logging
 from .base import PIPELINE_INIT_ARGS, ArgumentHandler, Pipeline
 
@@ -78,7 +79,14 @@ class ZeroShotClassificationPipeline(Pipeline):
         return -1
 
     def _parse_and_tokenize(
-        self, sequences, candidate_labels, hypothesis_template, padding=True, add_special_tokens=True, **kwargs
+        self,
+        sequences,
+        candidate_labels,
+        hypothesis_template,
+        padding=True,
+        add_special_tokens=True,
+        truncation=TruncationStrategy.ONLY_FIRST,
+        **kwargs
     ):
         """
         Parse arguments and tokenize only_first so that hypothesis (label) is not truncated
@@ -89,7 +97,7 @@ class ZeroShotClassificationPipeline(Pipeline):
             add_special_tokens=add_special_tokens,
             return_tensors=self.framework,
             padding=padding,
-            truncation="only_first",
+            truncation=truncation,
         )
 
         return inputs