Fix openvino model inference for transformers >= v4.29.3 (#329)

huggingface · May 31, 2023 · 2c8eb3a · 2c8eb3a
1 parent 7019728
commit 2c8eb3a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git 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
@@ -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: