Skip to content

GitLab

"tests/vscode:/vscode.git/clone" did not exist on "98c5e5da31dd70facf92970074be49501cd5e20b"
Unverified Commit 0a4b1068 authored by Sam Shleifer's avatar Sam Shleifer Committed by GitHub
Browse files

Speedup torch summarization tests (#3663)

parent 5aa8a278
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 8 deletions
tests/test_pipelines.py
View file @ 0a4b1068
...@@ -16,13 +16,11 @@ from .utils import require_tf, require_torch, slow ...@@ -16,13 +16,11 @@ from .utils import require_tf, require_torch, slow
QA_FINETUNED_MODELS = [ QA_FINETUNED_MODELS = [
(("bert-base-uncased", {"use_fast": False}), "bert-large-uncased-whole-word-masking-finetuned-squad", None), (("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), (("bert-base-cased", {"use_fast": False}), "distilbert-base-cased-distilled-squad", None),
] ]
TF_QA_FINETUNED_MODELS = [ TF_QA_FINETUNED_MODELS = [
(("bert-base-uncased", {"use_fast": False}), "bert-large-uncased-whole-word-masking-finetuned-squad", None), (("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), (("bert-base-cased", {"use_fast": False}), "distilbert-base-cased-distilled-squad", None),
] ]
...@@ -49,7 +47,6 @@ FEATURE_EXTRACT_FINETUNED_MODELS = { ...@@ -49,7 +47,6 @@ FEATURE_EXTRACT_FINETUNED_MODELS = {
} }
TF_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 # ('xlnet-base-cased', 'xlnet-base-cased', None), # Disabled for now as it crash for TF2
("distilbert-base-cased", "distilbert-base-cased", None), ("distilbert-base-cased", "distilbert-base-cased", None),
} }
...@@ -78,14 +75,17 @@ TF_FILL_MASK_FINETUNED_MODELS = [ ...@@ -78,14 +75,17 @@ TF_FILL_MASK_FINETUNED_MODELS = [
(("distilroberta-base", {"use_fast": False}), "distilroberta-base", None), (("distilroberta-base", {"use_fast": False}), "distilroberta-base", None),
] ]
SUMMARIZATION_FINETUNED_MODELS = {("bart-large-cnn", "bart-large-cnn"), ("t5-small", "t5-small")} SUMMARIZATION_FINETUNED_MODELS = {
TF_SUMMARIZATION_FINETUNED_MODELS = {("t5-small", "t5-small")} ("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 = { TRANSLATION_FINETUNED_MODELS = {
("t5-small", "t5-small", "translation_en_to_de"), ("patrickvonplaten/t5-tiny-random", "t5-small", "translation_en_to_de"),
("t5-small", "t5-small", "translation_en_to_ro"), ("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): class MonoColumnInputTestCase(unittest.TestCase):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment