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

Put smaller ALBERT model (#13028)

Browse Source
This commit is contained in:
Lysandre Debut
2021-08-06 18:41:33 +02:00
committed by GitHub
parent dc420b0eb1
commit 1bf38611a4
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_onnx_v2.py
View File

@@ -195,7 +195,7 @@ if is_torch_available():
) )
PYTORCH_EXPORT_DEFAULT_MODELS = { PYTORCH_EXPORT_DEFAULT_MODELS = {
("ALBERT", "albert-base-v2", AlbertModel, AlbertConfig, AlbertOnnxConfig), ("ALBERT", "hf-internal-testing/tiny-albert", AlbertModel, AlbertConfig, AlbertOnnxConfig),
("BART", "facebook/bart-base", BartModel, BartConfig, BartOnnxConfig), ("BART", "facebook/bart-base", BartModel, BartConfig, BartOnnxConfig),
("BERT", "bert-base-cased", BertModel, BertConfig, BertOnnxConfig), ("BERT", "bert-base-cased", BertModel, BertConfig, BertOnnxConfig),
("DistilBERT", "distilbert-base-cased", DistilBertModel, DistilBertConfig, DistilBertOnnxConfig), ("DistilBERT", "distilbert-base-cased", DistilBertModel, DistilBertConfig, DistilBertOnnxConfig),
@@ -229,7 +229,7 @@ class OnnxExportTestCaseV2(TestCase):
self.assertTrue(hasattr(onnx_config_class, "default")) self.assertTrue(hasattr(onnx_config_class, "default"))
tokenizer = AutoTokenizer.from_pretrained(model) tokenizer = AutoTokenizer.from_pretrained(model)
model = model_class(config_class()) model = model_class(config_class.from_pretrained(model))
onnx_config = onnx_config_class.default(model.config) onnx_config = onnx_config_class.default(model.config)
with NamedTemporaryFile("w") as output: with NamedTemporaryFile("w") as output: