SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Sam Shleifer
2020-05-18 12:23:21 -04:00
committed by GitHub
parent d9ece8233d
commit a699525d25
1 changed files with 26 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
tests/test_pipelines.py
View File

@@ -7,6 +7,8 @@ from transformers.pipelines import SUPPORTED_TASKS, DefaultArgumentHandler, Pipe
from .utils import require_tf, require_torch, slow 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"] NER_FINETUNED_MODELS = ["sshleifer/tiny-dbmdz-bert-large-cased-finetuned-conll03-english"]
# xlnet-base-cased disabled for now, since it crashes TF2 # 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}, {"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): class DefaultArgumentHandlerTestCase(unittest.TestCase):
@@ -119,10 +122,11 @@ class MonoColumnInputTestCase(unittest.TestCase):
invalid_inputs: List = [None], invalid_inputs: List = [None],
expected_multi_result: Optional[List] = None, expected_multi_result: Optional[List] = None,
expected_check_keys: Optional[List[str]] = None, expected_check_keys: Optional[List[str]] = None,
**kwargs,
): ):
self.assertIsNotNone(nlp) 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, list)
self.assertIsInstance(mono_result[0], (dict, list)) self.assertIsInstance(mono_result[0], (dict, list))
@@ -155,64 +159,56 @@ class MonoColumnInputTestCase(unittest.TestCase):
@require_torch @require_torch
def test_torch_ner(self): def test_torch_ner(self):
mandatory_keys = {"entity", "word", "score"} 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: for model_name in NER_FINETUNED_MODELS:
nlp = pipeline(task="ner", model=model_name, tokenizer=model_name) 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 @require_torch
def test_ner_grouped(self): def test_ner_grouped(self):
mandatory_keys = {"entity_group", "word", "score"} 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: for model_name in NER_FINETUNED_MODELS:
nlp = pipeline(task="ner", model=model_name, tokenizer=model_name, grouped_entities=True) 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 @require_tf
def test_tf_ner(self): def test_tf_ner(self):
mandatory_keys = {"entity", "word", "score"} 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: for model_name in NER_FINETUNED_MODELS:
nlp = pipeline(task="ner", model=model_name, tokenizer=model_name, framework="tf") 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 @require_tf
def test_tf_ner_grouped(self): def test_tf_ner_grouped(self):
mandatory_keys = {"entity_group", "word", "score"} 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: for model_name in NER_FINETUNED_MODELS:
nlp = pipeline(task="ner", model=model_name, tokenizer=model_name, framework="tf", grouped_entities=True) 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 @require_torch
def test_torch_sentiment_analysis(self): def test_torch_sentiment_analysis(self):
mandatory_keys = {"label", "score"} 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: for model_name in TEXT_CLASSIF_FINETUNED_MODELS:
nlp = pipeline(task="sentiment-analysis", model=model_name, tokenizer=model_name) 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 @require_tf
def test_tf_sentiment_analysis(self): def test_tf_sentiment_analysis(self):
mandatory_keys = {"label", "score"} 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: for model_name in TEXT_CLASSIF_FINETUNED_MODELS:
nlp = pipeline(task="sentiment-analysis", model=model_name, tokenizer=model_name, framework="tf") 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 @require_torch
def test_torch_feature_extraction(self): 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: for model_name in FEATURE_EXTRACT_FINETUNED_MODELS:
nlp = pipeline(task="feature-extraction", model=model_name, tokenizer=model_name) 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 @require_tf
def test_tf_feature_extraction(self): 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: for model_name in FEATURE_EXTRACT_FINETUNED_MODELS:
nlp = pipeline(task="feature-extraction", model=model_name, tokenizer=model_name, framework="tf") 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 @require_torch
def test_torch_fill_mask(self): def test_torch_fill_mask(self):
@@ -274,54 +270,55 @@ class MonoColumnInputTestCase(unittest.TestCase):
@require_torch @require_torch
def test_torch_summarization(self): def test_torch_summarization(self):
valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
invalid_inputs = [4, "<mask>"] invalid_inputs = [4, "<mask>"]
mandatory_keys = ["summary_text"] mandatory_keys = ["summary_text"]
for model in SUMMARIZATION_FINETUNED_MODELS: for model in SUMMARIZATION_FINETUNED_MODELS:
nlp = pipeline(task="summarization", model=model, tokenizer=model) 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 @require_tf
def test_tf_summarization(self): def test_tf_summarization(self):
valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
invalid_inputs = [4, "<mask>"] invalid_inputs = [4, "<mask>"]
mandatory_keys = ["summary_text"] mandatory_keys = ["summary_text"]
for model_name in TF_SUMMARIZATION_FINETUNED_MODELS: for model_name in TF_SUMMARIZATION_FINETUNED_MODELS:
nlp = pipeline(task="summarization", model=model_name, tokenizer=model_name, framework="tf",) 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 @require_torch
def test_torch_translation(self): def test_torch_translation(self):
valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
invalid_inputs = [4, "<mask>"] invalid_inputs = [4, "<mask>"]
mandatory_keys = ["translation_text"] mandatory_keys = ["translation_text"]
for model_name, task in TRANSLATION_FINETUNED_MODELS: for model_name, task in TRANSLATION_FINETUNED_MODELS:
nlp = pipeline(task=task, model=model_name, tokenizer=model_name) 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 @require_tf
@slow @slow
def test_tf_translation(self): def test_tf_translation(self):
valid_inputs = ["A string like this", ["list of strings entry 1", "list of strings v2"]]
invalid_inputs = [4, "<mask>"] invalid_inputs = [4, "<mask>"]
mandatory_keys = ["translation_text"] mandatory_keys = ["translation_text"]
for model, task in TF_TRANSLATION_FINETUNED_MODELS: for model, task in TF_TRANSLATION_FINETUNED_MODELS:
nlp = pipeline(task=task, model=model, tokenizer=model, framework="tf") 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 @require_torch
def test_torch_text_generation(self): 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: for model_name in TEXT_GENERATION_FINETUNED_MODELS:
nlp = pipeline(task="text-generation", model=model_name, tokenizer=model_name, framework="pt") 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 @require_tf
def test_tf_text_generation(self): 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: for model_name in TEXT_GENERATION_FINETUNED_MODELS:
nlp = pipeline(task="text-generation", model=model_name, tokenizer=model_name, framework="tf") 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"] QA_FINETUNED_MODELS = ["sshleifer/tiny-distilbert-base-cased-distilled-squad"]