uncomment modernbert

huggingface · Jan 13, 2025 · d1924f0 · d1924f0
1 parent c46de05
commit d1924f0
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 14 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -183,9 +183,9 @@ def inputs(self) -> Dict[str, Dict[int, str]]:
         return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis}
 
 
-# TODO: uncomment when transformers>=4.48.0
-# class ModernBertOnnxConfig(DistilBertOnnxConfig):
-#     pass
+class ModernBertOnnxConfig(DistilBertOnnxConfig):
+    MIN_TRANSFORMERS_VERSION = version.parse("4.48.0")
+    pass
 
 
 class MPNetOnnxConfig(DistilBertOnnxConfig):

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -893,15 +893,15 @@ class TasksManager:
             "image-classification",
             onnx="MobileNetV2OnnxConfig",
         ),
-        # "modernbert": supported_tasks_mapping(
-        #     "feature-extraction",
-        #     "fill-mask",
-        #     "text-classification",
-        #     "multiple-choice",
-        #     "token-classification",
-        #     "question-answering",
-        #     onnx="ModernBertOnnxConfig",
-        # ),
+        "modernbert": supported_tasks_mapping(
+            "feature-extraction",
+            "fill-mask",
+            "text-classification",
+            "multiple-choice",
+            "token-classification",
+            "question-answering",
+            onnx="ModernBertOnnxConfig",
+        ),
         "mpnet": supported_tasks_mapping(
             "feature-extraction",
             "fill-mask",

diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -125,7 +125,7 @@
     "mobilenet-v2": "hf-internal-testing/tiny-random-MobileNetV2Model",
     "mobilenet-v1": "google/mobilenet_v1_0.75_192",
     "mobilevit": "hf-internal-testing/tiny-random-mobilevit",
-    # "modernbert": "hf-internal-testing/tiny-random-ModernBertForMaskedLM",
+    "modernbert": "hf-internal-testing/tiny-random-ModernBertForMaskedLM",
     "mpnet": "hf-internal-testing/tiny-random-MPNetModel",
     "mpt": "hf-internal-testing/tiny-random-MptForCausalLM",
     "mt5": "lewtun/tiny-random-mt5",
@@ -269,7 +269,7 @@
     # "mobilenet_v1": "google/mobilenet_v1_0.75_192",
     # "mobilenet_v2": "google/mobilenet_v2_0.35_96",
     "mobilevit": "apple/mobilevit-small",
-    # "modernbert": "answerdotai/ModernBERT-base",
+    "modernbert": "answerdotai/ModernBERT-base",
     "mpt": "mosaicml/mpt-7b",
     "mt5": "lewtun/tiny-random-mt5",  # Not using google/mt5-small because it takes too much time for testing.
     "musicgen": "facebook/musicgen-small",