diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 4c5a727a18..564e98b55b 100644 --- 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 index 7cb5a31d2d..59c066ac38 100644 --- 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",