Checking output format + check raises ValueError (#8986)

src/transformers/pipelines.py

@@ -516,7 +516,7 @@ class Pipeline(_ScikitCompat):
         if framework is None:
             framework = get_framework(model)
 
-        self.task = task
+        self.task = task or getattr(self, "task", "")
         self.model = model
         self.tokenizer = tokenizer
         self.modelcard = modelcard
@@ -530,8 +530,8 @@ class Pipeline(_ScikitCompat):
 
         # Update config with task specific parameters
         task_specific_params = self.model.config.task_specific_params
-        if task_specific_params is not None and task in task_specific_params:
-            self.model.config.update(task_specific_params.get(task))
+        if task_specific_params is not None and self.task in task_specific_params:
+            self.model.config.update(task_specific_params.get(self.task))
 
     def save_pretrained(self, save_directory: str):
         """
@@ -926,6 +926,14 @@ class TextGenerationPipeline(Pipeline):
     r"""
         return_all_scores (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to return all prediction scores or just the one of the predicted class.
+        function_to_apply (:obj:`str`, `optional`, defaults to :obj:`"default"`):
+            The function to apply to the model outputs in order to retrieve the scores. Accepts four different values:
+
+            - :obj:`"default"`: if the model has a single label, will apply the sigmoid function on the output. If the
+              model has several labels, will apply the softmax function on the output.
+            - :obj:`"sigmoid"`: Applies the sigmoid function on the output.
+            - :obj:`"softmax"`: Applies the softmax function on the output.
+            - :obj:`"none"`: Does not apply any function on the output.
     """,
 )
 class TextClassificationPipeline(Pipeline):
@@ -945,7 +953,9 @@ class TextClassificationPipeline(Pipeline):
     <https://huggingface.co/models?filter=text-classification>`__.
     """
 
-    def __init__(self, return_all_scores: bool = False, **kwargs):
+    task = "text-classification"
+
+    def __init__(self, return_all_scores: bool = None, function_to_apply: str = None, **kwargs):
         super().__init__(**kwargs)
 
         self.check_model_type(
@@ -954,15 +964,33 @@ class TextClassificationPipeline(Pipeline):
             else MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
         )
 
-        self.return_all_scores = return_all_scores
+        if hasattr(self.model.config, "return_all_scores") and return_all_scores is None:
+            return_all_scores = self.model.config.return_all_scores
 
-    def __call__(self, *args, **kwargs):
+        if hasattr(self.model.config, "function_to_apply") and function_to_apply is None:
+            function_to_apply = self.model.config.function_to_apply
+
+        self.return_all_scores = return_all_scores if return_all_scores is not None else False
+        self.function_to_apply = function_to_apply if function_to_apply is not None else "default"
+
+    def __call__(self, *args, return_all_scores=None, function_to_apply=None, **kwargs):
         """
         Classify the text(s) given as inputs.
 
         Args:
             args (:obj:`str` or :obj:`List[str]`):
                 One or several texts (or one list of prompts) to classify.
+            return_all_scores (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether to return scores for all labels.
+            function_to_apply (:obj:`str`, `optional`, defaults to :obj:`"default"`):
+                The function to apply to the model outputs in order to retrieve the scores. Accepts four different
+                values:
+
+                - :obj:`"default"`: if the model has a single label, will apply the sigmoid function on the output. If
+                  the model has several labels, will apply the softmax function on the output.
+                - :obj:`"sigmoid"`: Applies the sigmoid function on the output.
+                - :obj:`"softmax"`: Applies the softmax function on the output.
+                - :obj:`"none"`: Does not apply any function on the output.
 
         Return:
             A list or a list of list of :obj:`dict`: Each result comes as list of dictionaries with the following keys:
@@ -974,11 +1002,30 @@ class TextClassificationPipeline(Pipeline):
         """
         outputs = super().__call__(*args, **kwargs)
 
-        if self.model.config.num_labels == 1:
-            scores = 1.0 / (1.0 + np.exp(-outputs))
+        return_all_scores = return_all_scores if return_all_scores is not None else self.return_all_scores
+        function_to_apply = function_to_apply if function_to_apply is not None else self.function_to_apply
+
+        def sigmoid(_outputs):
+            return 1.0 / (1.0 + np.exp(-_outputs))
+
+        def softmax(_outputs):
+            return np.exp(_outputs) / np.exp(_outputs).sum(-1, keepdims=True)
+
+        if function_to_apply == "default":
+            if self.model.config.num_labels == 1:
+                scores = sigmoid(outputs)
+            else:
+                scores = softmax(outputs)
+        elif function_to_apply == "sigmoid":
+            scores = sigmoid(outputs)
+        elif function_to_apply == "softmax":
+            scores = softmax(outputs)
+        elif function_to_apply.lower() == "none":
+            scores = outputs
         else:
-            scores = np.exp(outputs) / np.exp(outputs).sum(-1, keepdims=True)
-        if self.return_all_scores:
+            raise ValueError(f"Unrecognized `function_to_apply` argument: {function_to_apply}")
+
+        if return_all_scores:
             return [
                 [{"label": self.model.config.id2label[i], "score": score.item()} for i, score in enumerate(item)]
                 for item in scores
@@ -1040,6 +1087,8 @@ class ZeroShotClassificationPipeline(Pipeline):
     of available models on `huggingface.co/models <https://huggingface.co/models?search=nli>`__.
     """
 
+    task = "zero-shot"
+
     def __init__(self, args_parser=ZeroShotClassificationArgumentHandler(), *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._args_parser = args_parser
@@ -1172,6 +1221,8 @@ class FillMaskPipeline(Pipeline):
         This pipeline only works for inputs with exactly one token masked.
     """
 
+    task = "fill-mask"
+
     def __init__(
         self,
         model: Union["PreTrainedModel", "TFPreTrainedModel"],
@@ -1361,6 +1412,7 @@ class TokenClassificationPipeline(Pipeline):
     <https://huggingface.co/models?filter=token-classification>`__.
     """
 
+    task = "token-classification"
     default_input_names = "sequences"
 
     def __init__(
@@ -1667,6 +1719,7 @@ class QuestionAnsweringPipeline(Pipeline):
     <https://huggingface.co/models?filter=question-answering>`__.
     """
 
+    task = "question-answering"
     default_input_names = "question,context"
 
     def __init__(
@@ -2067,6 +2120,8 @@ class SummarizationPipeline(Pipeline):
         summarizer("Sam Shleifer writes the best docstring examples in the whole world.", min_length=5, max_length=20)
     """
 
+    task = "summarization"
+
     def __init__(self, *args, **kwargs):
         kwargs.update(task="summarization")
         super().__init__(*args, **kwargs)
@@ -2188,6 +2243,8 @@ class TranslationPipeline(Pipeline):
         en_fr_translator("How old are you?")
     """
 
+    task = "translation"
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -2297,6 +2354,8 @@ class Text2TextGenerationPipeline(Pipeline):
         text2text_generator("question: What is 42 ? context: 42 is the answer to life, the universe and everything")
     """
 
+    task = "text2text-generation"
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -2522,6 +2581,8 @@ class ConversationalPipeline(Pipeline):
         conversational_pipeline([conversation_1, conversation_2])
     """
 
+    task = "conversational"
+
     def __init__(self, min_length_for_response=32, *args, **kwargs):
         super().__init__(*args, **kwargs)
 

tests/test_pipelines_common.py

@@ -197,6 +197,7 @@ class MonoInputPipelineCommonMixin(CustomInputPipelineCommonMixin):
 
     def _test_pipeline(self, nlp: Pipeline):
         self.assertIsNotNone(nlp)
+        self.assertIsNotNone(nlp.task)
 
         mono_result = nlp(self.valid_inputs[0], **self.pipeline_running_kwargs)
         self.assertIsInstance(mono_result, list)

tests/test_pipelines_sentiment_analysis.py

@@ -1,12 +1,117 @@
 import unittest
 
+import numpy as np
+
+from transformers import AutoTokenizer, DistilBertConfig, DistilBertForSequenceClassification, pipeline
+from transformers.testing_utils import slow
+
 from .test_pipelines_common import MonoInputPipelineCommonMixin
 
 
+VALID_INPUTS = ["I really disagree with what you've said.", ["I love you."]]
+
+
 class SentimentAnalysisPipelineTests(MonoInputPipelineCommonMixin, unittest.TestCase):
     pipeline_task = "sentiment-analysis"
-    small_models = [
-        "sshleifer/tiny-distilbert-base-uncased-finetuned-sst-2-english"
-    ]  # Default model - Models tested without the @slow decorator
+    small_models = ["distilbert-base-cased"]  # Default model - Models tested without the @slow decorator
     large_models = [None]  # Models tested with the @slow decorator
     mandatory_keys = {"label", "score"}  # Keys which should be in the output
+
+    @slow
+    def test_function_to_apply(self):
+        for model_name in self.small_models:
+            string_input, string_list_input = VALID_INPUTS
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+
+            model = DistilBertForSequenceClassification(DistilBertConfig.from_pretrained(model_name))
+            model.eval()
+            classifier = pipeline(task="sentiment-analysis", model=model, tokenizer=tokenizer)
+
+            def check_output(output):
+                # This model does not have a sequence-classification head, so results are random
+                self.assertTrue(isinstance(output, list))
+                self.assertEqual(len(output), 1)
+                d = string_output[0]
+                self.assertEqual(set(d.keys()), {"label", "score"})
+                self.assertEqual(type(d["label"]), str)
+                self.assertEqual(type(d["score"]), float)
+
+            def pipeline_call_argument(argument=None):
+                _string_output = classifier(string_input, function_to_apply=argument)
+                _string_list_output = classifier(string_list_input, function_to_apply=argument)
+                check_output(_string_output)
+                return _string_output, _string_list_output
+
+            def pipeline_init_argument(argument=None):
+                _classifier = pipeline(
+                    task="sentiment-analysis", model=model, tokenizer=tokenizer, function_to_apply=argument
+                )
+                _string_output = _classifier(string_input)
+                _string_list_output = _classifier(string_list_input)
+                check_output(_string_output)
+                return _string_output, _string_list_output
+
+            def pipeline_model_argument(argument=None):
+                model.config.task_specific_params = {"sentiment-analysis": {"function_to_apply": argument}}
+                _classifier = pipeline(task="sentiment-analysis", model=model, tokenizer=tokenizer)
+                _string_output = _classifier(string_input)
+                _string_list_output = _classifier(string_list_input)
+                check_output(_string_output)
+                return _string_output, _string_list_output
+
+            string_output = classifier(string_input)
+            string_list_output = classifier(string_list_input)
+
+            string_output_default = pipeline_call_argument("default")
+            string_output_sigmoid = pipeline_call_argument("sigmoid")
+            string_output_softmax = pipeline_call_argument("softmax")
+            string_output_none = pipeline_call_argument("none")
+
+            string_output_init_default = pipeline_init_argument("default")
+            string_output_init_sigmoid = pipeline_init_argument("sigmoid")
+            string_output_init_softmax = pipeline_init_argument("softmax")
+            string_output_init_none = pipeline_init_argument("none")
+
+            string_output_model_default = pipeline_model_argument("default")
+            string_output_model_sigmoid = pipeline_model_argument("sigmoid")
+            string_output_model_softmax = pipeline_model_argument("softmax")
+            string_output_model_none = pipeline_model_argument("none")
+
+            should_be_equal = [
+                (
+                    (string_output, string_list_output),
+                    string_output_default,
+                    string_output_init_default,
+                    string_output_model_default,
+                ),
+                (string_output_sigmoid, string_output_init_sigmoid, string_output_model_sigmoid),
+                (string_output_softmax, string_output_init_softmax, string_output_model_softmax),
+                (string_output_none, string_output_init_none, string_output_model_none),
+            ]
+
+            for tuples_containing_equal_values in should_be_equal:
+                # Retrieve each tuple from the list
+                for tuple_value_0 in tuples_containing_equal_values:
+                    for tuple_value_1 in tuples_containing_equal_values:
+                        if tuple_value_0 is not tuple_value_1:
+                            # Compare each tuple value with all the others, as long as they're not the same object
+                            for pipeline_output_0, pipeline_output_1 in zip(tuple_value_0, tuple_value_1):
+                                # Compare all outputs (call, init, model argument)
+                                for example_result_0, example_result_1 in zip(pipeline_output_0, pipeline_output_1):
+                                    # Iterate through the results
+                                    self.assertTrue(
+                                        np.allclose(example_result_0["score"], example_result_1["score"], atol=1e-6)
+                                    )
+
+    @slow
+    def test_function_to_apply_error(self):
+        for model_name in self.small_models:
+            string_input, _ = VALID_INPUTS
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+
+            model = DistilBertForSequenceClassification(DistilBertConfig.from_pretrained(model_name))
+            model.eval()
+            classifier = pipeline(task="sentiment-analysis", model=model, tokenizer=tokenizer)
+
+            with self.assertRaises(ValueError):
+                classifier(string_input, function_to_apply="logits")