Mrbean/codegen onnx (#17903)

This commit is contained in:
mrbean
2022-06-28 08:57:53 -04:00
committed by GitHub
parent 76d13de5ae
commit b424f0b4a3
2 changed files with 6 additions and 0 deletions

View File

@@ -188,6 +188,11 @@ class FeaturesManager:
"question-answering", "question-answering",
onnx_config_cls="models.camembert.CamembertOnnxConfig", onnx_config_cls="models.camembert.CamembertOnnxConfig",
), ),
"codegen": supported_features_mapping(
"default",
"causal-lm",
onnx_config_cls="models.codegen.CodeGenOnnxConfig",
),
"convbert": supported_features_mapping( "convbert": supported_features_mapping(
"default", "default",
"masked-lm", "masked-lm",

View File

@@ -180,6 +180,7 @@ PYTORCH_EXPORT_MODELS = {
("ibert", "kssteven/ibert-roberta-base"), ("ibert", "kssteven/ibert-roberta-base"),
("camembert", "camembert-base"), ("camembert", "camembert-base"),
("convbert", "YituTech/conv-bert-base"), ("convbert", "YituTech/conv-bert-base"),
("codegen", "Salesforce/codegen-350M-multi"),
("deberta", "microsoft/deberta-base"), ("deberta", "microsoft/deberta-base"),
("deberta-v2", "microsoft/deberta-v2-xlarge"), ("deberta-v2", "microsoft/deberta-v2-xlarge"),
("convnext", "facebook/convnext-tiny-224"), ("convnext", "facebook/convnext-tiny-224"),