[test_pipelines] Mark tests > 10s @slow, small speedups (#4421)

tests/test_pipelines.py

@@ -7,6 +7,8 @@ from transformers.pipelines import SUPPORTED_TASKS, DefaultArgumentHandler, Pipe
 from .utils import require_tf, require_torch, slow
 
 
+VALID_INPUTS = ["A simple string", ["list of strings"]]
+
 NER_FINETUNED_MODELS = ["sshleifer/tiny-dbmdz-bert-large-cased-finetuned-conll03-english"]
 
 # xlnet-base-cased disabled for now, since it crashes TF2
@@ -36,6 +38,7 @@ expected_fill_mask_result = [
         {"sequence": "<s> The largest city in France is Lyon</s>", "score": 0.21112334728240967, "token": 12790},
     ],
 ]
+SUMMARIZATION_KWARGS = dict(num_beams=2, min_length=2, max_length=5)
 
 
 class DefaultArgumentHandlerTestCase(unittest.TestCase):
@@ -119,10 +122,11 @@ class MonoColumnInputTestCase(unittest.TestCase):
         invalid_inputs: List = [None],
         expected_multi_result: Optional[List] = None,
         expected_check_keys: Optional[List[str]] = None,
+        **kwargs,
     ):
         self.assertIsNotNone(nlp)
 
-        mono_result = nlp(valid_inputs[0])
+        mono_result = nlp(valid_inputs[0], **kwargs)
         self.assertIsInstance(mono_result, list)
         self.assertIsInstance(mono_result[0], (dict, list))
 
@@ -155,64 +159,56 @@ class MonoColumnInputTestCase(unittest.TestCase):
     @require_torch
     def test_torch_ner(self):
         mandatory_keys = {"entity", "word", "score"}
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in NER_FINETUNED_MODELS:
             nlp = pipeline(task="ner", model=model_name, tokenizer=model_name)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys)
 
     @require_torch
     def test_ner_grouped(self):
         mandatory_keys = {"entity_group", "word", "score"}
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in NER_FINETUNED_MODELS:
             nlp = pipeline(task="ner", model=model_name, tokenizer=model_name, grouped_entities=True)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys)
 
     @require_tf
     def test_tf_ner(self):
         mandatory_keys = {"entity", "word", "score"}
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in NER_FINETUNED_MODELS:
             nlp = pipeline(task="ner", model=model_name, tokenizer=model_name, framework="tf")
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys)
 
     @require_tf
     def test_tf_ner_grouped(self):
         mandatory_keys = {"entity_group", "word", "score"}
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in NER_FINETUNED_MODELS:
             nlp = pipeline(task="ner", model=model_name, tokenizer=model_name, framework="tf", grouped_entities=True)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys)
 
     @require_torch
     def test_torch_sentiment_analysis(self):
         mandatory_keys = {"label", "score"}
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in TEXT_CLASSIF_FINETUNED_MODELS:
             nlp = pipeline(task="sentiment-analysis", model=model_name, tokenizer=model_name)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys)
 
     @require_tf
     def test_tf_sentiment_analysis(self):
         mandatory_keys = {"label", "score"}
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in TEXT_CLASSIF_FINETUNED_MODELS:
             nlp = pipeline(task="sentiment-analysis", model=model_name, tokenizer=model_name, framework="tf")
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys)
 
     @require_torch
     def test_torch_feature_extraction(self):
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in FEATURE_EXTRACT_FINETUNED_MODELS:
             nlp = pipeline(task="feature-extraction", model=model_name, tokenizer=model_name)
-            self._test_mono_column_pipeline(nlp, valid_inputs, {})
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, {})
 
     @require_tf
     def test_tf_feature_extraction(self):
-        valid_inputs = ["HuggingFace is solving NLP one commit at a time.", "HuggingFace is based in New-York & Paris"]
         for model_name in FEATURE_EXTRACT_FINETUNED_MODELS:
             nlp = pipeline(task="feature-extraction", model=model_name, tokenizer=model_name, framework="tf")
-            self._test_mono_column_pipeline(nlp, valid_inputs, {})
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, {})
 
     @require_torch
     def test_torch_fill_mask(self):
@@ -274,54 +270,55 @@ class MonoColumnInputTestCase(unittest.TestCase):
 
     @require_torch
     def test_torch_summarization(self):
-        valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
         invalid_inputs = [4, "<mask>"]
         mandatory_keys = ["summary_text"]
         for model in SUMMARIZATION_FINETUNED_MODELS:
             nlp = pipeline(task="summarization", model=model, tokenizer=model)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys, invalid_inputs=invalid_inputs)
+            self._test_mono_column_pipeline(
+                nlp, VALID_INPUTS, mandatory_keys, invalid_inputs=invalid_inputs, **SUMMARIZATION_KWARGS
+            )
 
+    @slow
     @require_tf
     def test_tf_summarization(self):
-        valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
         invalid_inputs = [4, "<mask>"]
         mandatory_keys = ["summary_text"]
         for model_name in TF_SUMMARIZATION_FINETUNED_MODELS:
             nlp = pipeline(task="summarization", model=model_name, tokenizer=model_name, framework="tf",)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys, invalid_inputs=invalid_inputs)
+            self._test_mono_column_pipeline(
+                nlp, VALID_INPUTS, mandatory_keys, invalid_inputs=invalid_inputs, **SUMMARIZATION_KWARGS
+            )
 
     @require_torch
     def test_torch_translation(self):
-        valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
         invalid_inputs = [4, "<mask>"]
         mandatory_keys = ["translation_text"]
         for model_name, task in TRANSLATION_FINETUNED_MODELS:
             nlp = pipeline(task=task, model=model_name, tokenizer=model_name)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys, invalid_inputs)
+            self._test_mono_column_pipeline(
+                nlp, VALID_INPUTS, mandatory_keys, invalid_inputs,
+            )
 
     @require_tf
     @slow
     def test_tf_translation(self):
-        valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
         invalid_inputs = [4, "<mask>"]
         mandatory_keys = ["translation_text"]
         for model, task in TF_TRANSLATION_FINETUNED_MODELS:
             nlp = pipeline(task=task, model=model, tokenizer=model, framework="tf")
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys, invalid_inputs=invalid_inputs)
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, mandatory_keys, invalid_inputs=invalid_inputs)
 
     @require_torch
     def test_torch_text_generation(self):
-        valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
         for model_name in TEXT_GENERATION_FINETUNED_MODELS:
             nlp = pipeline(task="text-generation", model=model_name, tokenizer=model_name, framework="pt")
-            self._test_mono_column_pipeline(nlp, valid_inputs, {})
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, {})
 
     @require_tf
     def test_tf_text_generation(self):
-        valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
         for model_name in TEXT_GENERATION_FINETUNED_MODELS:
             nlp = pipeline(task="text-generation", model=model_name, tokenizer=model_name, framework="tf")
-            self._test_mono_column_pipeline(nlp, valid_inputs, {})
+            self._test_mono_column_pipeline(nlp, VALID_INPUTS, {})
 
 
 QA_FINETUNED_MODELS = ["sshleifer/tiny-distilbert-base-cased-distilled-squad"]