Adds MPNet to NormalizedConfig and ORTConfigManager (#1471)

Adds mpnet to ORTConfigManager and NormalizedConfigManager Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com>
huggingface · Feb 29, 2024 · 64f795d · 64f795d
1 parent bb21ae7
commit 64f795d
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 1 deletion.
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
@@ -122,6 +122,7 @@ class ORTConfigManager:
         "marian": "bart",
         "mbart": "bart",
         "mistral": "gpt2",
+        "mpnet": "bert",
         "mt5": "bart",
         "m2m-100": "bart",
         "nystromformer": "bert",

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -241,6 +241,7 @@ class NormalizedConfigManager:
         "mbart": BartLikeNormalizedTextConfig,
         "mistral": NormalizedTextConfigWithGQA,
         "mixtral": NormalizedTextConfigWithGQA,
+        "mpnet": NormalizedTextConfig,
         "mpt": MPTNormalizedTextConfig,
         "mt5": T5LikeNormalizedTextConfig,
         "m2m-100": BartLikeNormalizedTextConfig,

diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -1435,6 +1435,7 @@ class ORTModelForMaskedLMIntegrationTest(ORTModelTestMixin):
         "flaubert",
         "ibert",
         "mobilebert",
+        "mpnet",
         "perceiver_text",
         "roberta",
         "roformer",
@@ -1976,7 +1977,16 @@ def test_compare_to_io_binding(self, model_arch):
 
 
 class ORTModelForFeatureExtractionIntegrationTest(ORTModelTestMixin):
-    SUPPORTED_ARCHITECTURES = ["albert", "bert", "camembert", "distilbert", "electra", "roberta", "xlm_roberta"]
+    SUPPORTED_ARCHITECTURES = [
+        "albert",
+        "bert",
+        "camembert",
+        "distilbert",
+        "electra",
+        "mpnet",
+        "roberta",
+        "xlm_roberta",
+    ]
 
     FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES}
     ORTMODEL_CLASS = ORTModelForFeatureExtraction

diff --git a/tests/onnxruntime/utils_onnxruntime_tests.py b/tests/onnxruntime/utils_onnxruntime_tests.py
@@ -120,6 +120,7 @@
     "mobilenet_v1": "google/mobilenet_v1_0.75_192",
     "mobilenet_v2": "hf-internal-testing/tiny-random-MobileNetV2Model",
     "mobilevit": "hf-internal-testing/tiny-random-mobilevit",
+    "mpnet": "hf-internal-testing/tiny-random-MPNetModel",
     "mpt": "hf-internal-testing/tiny-random-MptForCausalLM",
     "mt5": "lewtun/tiny-random-mt5",
     "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel",