fix(onnx/granite): Use LlamaOnnxConfig as the base for GraniteOnnxConfig

Branch: OnnxGranite Signed-off-by: Gabe Goodhart <[email protected]>
huggingface · Oct 22, 2024 · 3a96b22 · 3a96b22
1 parent e4f8707
commit 3a96b22
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -290,14 +290,8 @@ class GemmaOnnxConfig(LlamaOnnxConfig):
     pass
 
 
-class GraniteOnnxConfig(TextDecoderOnnxConfig):
-    # GG TODO: Bump past 4.44.2 once the next release is out
-    MIN_TRANSFORMERS_VERSION = version.parse("4.44.2")
-    DEFAULT_ONNX_OPSET = 14  # Granite follows Llama's default
-
-    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator)
-    DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator
-    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
+class GraniteOnnxConfig(LlamaOnnxConfig):
+    MIN_TRANSFORMERS_VERSION = version.parse("4.45.0")
 
 
 class PhiOnnxConfig(TextDecoderWithPositionIdsOnnxConfig):