diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index a13edc8f13..b97cda3834 100644 --- a/tests/test_pipelines.py +++ b/tests/test_pipelines.py @@ -16,13 +16,11 @@ from .utils import require_tf, require_torch, slow QA_FINETUNED_MODELS = [ (("bert-base-uncased", {"use_fast": False}), "bert-large-uncased-whole-word-masking-finetuned-squad", None), - (("bert-base-cased", {"use_fast": False}), "bert-large-cased-whole-word-masking-finetuned-squad", None), (("bert-base-cased", {"use_fast": False}), "distilbert-base-cased-distilled-squad", None), ] TF_QA_FINETUNED_MODELS = [ (("bert-base-uncased", {"use_fast": False}), "bert-large-uncased-whole-word-masking-finetuned-squad", None), - (("bert-base-cased", {"use_fast": False}), "bert-large-cased-whole-word-masking-finetuned-squad", None), (("bert-base-cased", {"use_fast": False}), "distilbert-base-cased-distilled-squad", None), ] @@ -49,7 +47,6 @@ FEATURE_EXTRACT_FINETUNED_MODELS = { } TF_FEATURE_EXTRACT_FINETUNED_MODELS = { - ("bert-base-cased", "bert-base-cased", None), # ('xlnet-base-cased', 'xlnet-base-cased', None), # Disabled for now as it crash for TF2 ("distilbert-base-cased", "distilbert-base-cased", None), } @@ -78,14 +75,17 @@ TF_FILL_MASK_FINETUNED_MODELS = [ (("distilroberta-base", {"use_fast": False}), "distilroberta-base", None), ] -SUMMARIZATION_FINETUNED_MODELS = {("bart-large-cnn", "bart-large-cnn"), ("t5-small", "t5-small")} -TF_SUMMARIZATION_FINETUNED_MODELS = {("t5-small", "t5-small")} +SUMMARIZATION_FINETUNED_MODELS = { + ("sshleifer/bart-tiny-random", "bart-large-cnn"), + ("patrickvonplaten/t5-tiny-random", "t5-small"), +} +TF_SUMMARIZATION_FINETUNED_MODELS = {("patrickvonplaten/t5-tiny-random", "t5-small")} TRANSLATION_FINETUNED_MODELS = { - ("t5-small", "t5-small", "translation_en_to_de"), - ("t5-small", "t5-small", "translation_en_to_ro"), + ("patrickvonplaten/t5-tiny-random", "t5-small", "translation_en_to_de"), + ("patrickvonplaten/t5-tiny-random", "t5-small", "translation_en_to_ro"), } -TF_TRANSLATION_FINETUNED_MODELS = {("t5-small", "t5-small", "translation_en_to_fr")} +TF_TRANSLATION_FINETUNED_MODELS = {("patrickvonplaten/t5-tiny-random", "t5-small", "translation_en_to_fr")} class MonoColumnInputTestCase(unittest.TestCase):