Add ONNX export support for modernbert

huggingface · Dec 19, 2024 · 282177a · 282177a
1 parent d33710d
commit 282177a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -179,6 +179,10 @@ def inputs(self) -> Dict[str, Dict[int, str]]:
         return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis}
 
 
+class ModernBertOnnxConfig(DistilBertOnnxConfig):
+    pass
+
+
 class MPNetOnnxConfig(DistilBertOnnxConfig):
     DEFAULT_ONNX_OPSET = 12  # For lower opsets, results in: Type 'tensor(int64)' of input parameter (/0/auto_model/encoder/Add_1_output_0) of operator (Min) in node (/0/auto_model/encoder/Min) is invalid.
 

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -891,6 +891,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",
+        ),
         "mpnet": supported_tasks_mapping(
             "feature-extraction",
             "fill-mask",