bwang0911 commited on
Commit
bc83a13
1 Parent(s): f713825

Update configuration_bert.py

Browse files
Files changed (1) hide show
  1. configuration_bert.py +3 -3
configuration_bert.py CHANGED
@@ -19,7 +19,7 @@ from collections import OrderedDict
19
  from typing import Mapping
20
 
21
  from transformers.configuration_utils import PretrainedConfig
22
- from transformers.onnx import OnnxConfig
23
  from transformers.utils import logging
24
 
25
 
@@ -152,7 +152,8 @@ class JinaBertConfig(PretrainedConfig):
152
  self.emb_pooler = emb_pooler
153
  self.attn_implementation = attn_implementation
154
 
155
- class JinaBertOnnxConfig(OnnxConfig):
 
156
  @property
157
  def inputs(self) -> Mapping[str, Mapping[int, str]]:
158
  if self.task == "multiple-choice":
@@ -163,6 +164,5 @@ class JinaBertOnnxConfig(OnnxConfig):
163
  [
164
  ("input_ids", dynamic_axis),
165
  ("attention_mask", dynamic_axis),
166
- ("token_type_ids", dynamic_axis),
167
  ]
168
  )
 
19
  from typing import Mapping
20
 
21
  from transformers.configuration_utils import PretrainedConfig
22
+ from optimum.exporters.onnx.model_configs import OnnxConfig, BertOnnxConfig
23
  from transformers.utils import logging
24
 
25
 
 
152
  self.emb_pooler = emb_pooler
153
  self.attn_implementation = attn_implementation
154
 
155
+ class JinaBertOnnxConfig(BertOnnxConfig):
156
+
157
  @property
158
  def inputs(self) -> Mapping[str, Mapping[int, str]]:
159
  if self.task == "multiple-choice":
 
164
  [
165
  ("input_ids", dynamic_axis),
166
  ("attention_mask", dynamic_axis),
 
167
  ]
168
  )