Updated configs for MPT models.

huggingface · Jun 28, 2023 · c4a1866 · c4a1866
1 parent 94690e5
commit c4a1866
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -228,6 +228,7 @@ def generate(self, input_name: str, framework: str = "pt"):
             for _ in range(self.num_layers)
         ]
 
+
 class MPTOnnxConfig(TextDecoderOnnxConfig):
     DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MPTPastKeyValuesGenerator)
     DEFAULT_ONNX_OPSET = 14

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -654,6 +654,13 @@ class TasksManager:
             onnx="MPNetOnnxConfig",
             tflite="MPNetTFLiteConfig",
         ),
+        "mpt": supported_tasks_mapping(
+            "feature-extraction",
+            "feature-extraction-with-past",
+            "text-generation",
+            "text-generation-with-past",
+            onnx="MPTOnnxConfig",
+        ),
         "mt5": supported_tasks_mapping(
             "feature-extraction",
             "feature-extraction-with-past",

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -222,6 +222,7 @@ class NormalizedConfigManager:
         "marian": BartLikeNormalizedTextConfig,
         "mbart": BartLikeNormalizedTextConfig,
         "mt5": T5LikeNormalizedTextConfig,
+        "mpt": NormalizedTextConfig.with_args(hidden_size="d_model", num_attention_heads="n_heads", num_layers="n_layers"),
         "m2m_100": BartLikeNormalizedTextConfig,
         "nystromformer": NormalizedTextConfig,
         "opt": NormalizedTextConfig,