skip optimmum if not installed (#2)

- skip optimmum if not installed (e7968d5b22061a5a5a5f996b52bf284520722c94)


Co-authored-by: Joan Fontanals Martínez <[email protected]>

configuration_bert.py

@@ -17,11 +17,18 @@
 """ BERT model configuration"""
 from collections import OrderedDict
 from typing import Mapping
+import warnings
 
 from transformers.configuration_utils import PretrainedConfig
-from optimum.exporters.onnx.model_configs import OnnxConfig, BertOnnxConfig
 from transformers.utils import logging
 
+try:
+    from optimum.exporters.onnx.model_configs import BertOnnxConfig
+    OPTIMUM_INSTALLED = True
+except ImportError:
+    warnings.warn("optimum is not installed. To use OnnxConfig and BertOnnxConfig, make sure that `optimum` package is installed")
+    OPTIMUM_INSTALLED = False
+
 
 logger = logging.get_logger(__name__)
 
@@ -152,17 +159,19 @@ class JinaBertConfig(PretrainedConfig):
         self.emb_pooler = emb_pooler
         self.attn_implementation = attn_implementation
 
-class JinaBertOnnxConfig(BertOnnxConfig):
-    
-    @property
-    def inputs(self) -> Mapping[str, Mapping[int, str]]:
-        if self.task == "multiple-choice":
-            dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
-        else:
-            dynamic_axis = {0: "batch", 1: "sequence"}
-        return OrderedDict(
-            [
-                ("input_ids", dynamic_axis),
-                ("attention_mask", dynamic_axis),
-            ]
-        )
+if OPTIMUM_INSTALLED:
+
+    class JinaBertOnnxConfig(BertOnnxConfig):
+
+        @property
+        def inputs(self) -> Mapping[str, Mapping[int, str]]:
+            if self.task == "multiple-choice":
+                dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
+            else:
+                dynamic_axis = {0: "batch", 1: "sequence"}
+            return OrderedDict(
+                [
+                    ("input_ids", dynamic_axis),
+                    ("attention_mask", dynamic_axis),
+                ]
+            )