Skip to content

GitLab

"...git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "1670be4bdec19d5a8893f943bf78a8d9b3dc8911"
Unverified Commit 1bf38611 authored by Lysandre Debut's avatar Lysandre Debut Committed by GitHub
Browse files

Put smaller ALBERT model (#13028)

parent dc420b0e
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
tests/test_onnx_v2.py
View file @ 1bf38611
...@@ -195,7 +195,7 @@ if is_torch_available(): ...@@ -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): ...@@ -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:
......
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