diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 9dbcacb7f..5276ade33 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -69,6 +69,7 @@ GptNeoxJapaneseModelPatcher, GptNeoxModelPatcher, IBertModelPatcher, + InputEmbeddingPatcher, InternLM2Patcher, InternLMModelPatcher, InternVLChatImageEmbeddingModelPatcher, @@ -1264,6 +1265,12 @@ def rename_ambiguous_inputs(self, inputs): model_inputs["input"] = inputs["input_ids"] return model_inputs + def patch_model_for_export( + self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None + ) -> "ModelPatcher": + # making 16bit tracable overrides embeedings input signature these changes required to prevent this issue + return InputEmbeddingPatcher(self, model, model_kwargs) + class LlavaConfigBehavior(str, enum.Enum): LANGUAGE = "language" diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 8507d94fe..dbbfb5662 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -2991,3 +2991,24 @@ def __init__( def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) self._model.forward = self._model.__orig_forward + + +class InputEmbeddingPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: Union["PreTrainedModel", "TFPreTrainedModel"], + model_kwargs: Dict[str, Any], + ): + model.__orig_forward = model.forward + + def forward(self, input): + return self.__orig_forward(input) + + model.forward = types.MethodType(forward, model) + + super().__init__(config, model, model_kwargs) + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + self._model.forward = self._model.__orig_forward