Skip to content

GitLab

Unverified Commit b424f0b4 authored by mrbean's avatar mrbean Committed by GitHub
Browse files

Mrbean/codegen onnx (#17903)

parent 76d13de5
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 0 deletions
src/transformers/onnx/features.py
View file @ b424f0b4
...@@ -188,6 +188,11 @@ class FeaturesManager: ...@@ -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",
......
tests/onnx/test_onnx_v2.py
View file @ b424f0b4
...@@ -180,6 +180,7 @@ PYTORCH_EXPORT_MODELS = { ...@@ -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"),
......
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