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

[tests] make pipelines tests faster with smaller models (#4238)

Browse Source 
covers torch and tf. Also fixes a failing @slow test
This commit is contained in:
Sam Shleifer
2020-05-14 13:36:02 -04:00
committed by GitHub
parent 448c467256
commit 7822cd38a0
2 changed files with 140 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/transformers/pipelines.py
View File

@@ -1513,7 +1513,7 @@ class TranslationPipeline(Pipeline):
return results
# Register all the supported task here
# Register all the supported tasks here
SUPPORTED_TASKS = {
"feature-extraction": {
"impl": FeatureExtractionPipeline,
@@ -1576,9 +1576,9 @@ SUPPORTED_TASKS = {
"tf": TFAutoModelWithLMHead if is_tf_available() else None,
"pt": AutoModelWithLMHead if is_torch_available() else None,
"default": {
"model": {"pt": "bart-large-cnn", "tf": None},
"model": {"pt": "bart-large-cnn", "tf": "t5-small"},
"config": None,
"tokenizer": ("bart-large-cnn", {"use_fast": False}),
"tokenizer": {"pt": ("bart-large-cnn", {"use_fast": False}), "tf": "t5-small"},
},
},
"translation_en_to_fr": {