From 2c8eb3aa1b8f94a52182bcf3177b38d2a4c0eb20 Mon Sep 17 00:00:00 2001
From: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>
Date: Wed, 31 May 2023 18:46:48 +0200
Subject: [PATCH] Fix openvino model inference for transformers >= v4.29.3
 (#329)

---
 optimum/intel/openvino/modeling_base.py | 7 ++++++-
 optimum/intel/openvino/trainer.py       | 2 ++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py
index ad0ffa3524..960ec4a4e9 100644
--- a/optimum/intel/openvino/modeling_base.py
+++ b/optimum/intel/openvino/modeling_base.py
@@ -55,12 +55,17 @@
 }
 
 
+# workaround to enable compatibility between openvino models and transformers pipelines
+class PreTrainedModel(OptimizedModel):
+    pass
+
+
 @add_start_docstrings(
     """
     Base OVModel class.
     """,
 )
-class OVBaseModel(OptimizedModel):
+class OVBaseModel(PreTrainedModel):
     _AUTOMODELS_TO_TASKS = {cls_name: task for task, cls_name in TasksManager._TASKS_TO_AUTOMODELS.items()}
     auto_model_class = None
     export_feature = None
diff --git a/optimum/intel/openvino/trainer.py b/optimum/intel/openvino/trainer.py
index e16bf0f58d..e51473028d 100644
--- a/optimum/intel/openvino/trainer.py
+++ b/optimum/intel/openvino/trainer.py
@@ -194,6 +194,8 @@ def __init__(
 
             self.compression_controller, self.model = create_compressed_model(self.model, nncf_config)
             self.model_wrapped = self.model
+            # TODO : To deprecate once support transformers > 4.30.0
+            self.deepspeed = None
 
     def _set_signature_columns_if_needed(self):
         if self._signature_columns is None: