From a699525d25142eb01741cc336022fdad216f343c Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Mon, 18 May 2020 12:23:21 -0400 Subject: [PATCH] [test_pipelines] Mark tests > 10s @slow, small speedups (#4421) --- tests/test_pipelines.py | 55 +++++++++++++++++++---------------------- 1 file changed, 26 insertions(+), 29 deletions(-) diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index 66de91e5df..02eede6921 100644 --- a/tests/test_pipelines.py +++ b/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": " The largest city in France is Lyon", "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, ""] 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, ""] 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, ""] 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, ""] 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"]