Fix for bloom family

huggingface · Jan 8, 2024 · 3bf4272 · 3bf4272
1 parent 49a7f8f
commit 3bf4272
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/optimum/exporters/openvino/stateful.py b/optimum/exporters/openvino/stateful.py
@@ -22,7 +22,11 @@
 from packaging import version
 import openvino as ov
 from openvino.runtime import opset13
+<<<<<<< HEAD
 from optimum.intel.utils.import_utils import _openvino_version, is_openvino_version
+=======
+from optimum.intel.utils.import_utils import is_openvino_version
+>>>>>>> Fix for bloom family
 from optimum.utils.normalized_config import NormalizedConfigManager
 
 
@@ -224,8 +228,13 @@ def patch_stateful(config: PretrainedConfig, ov_model: ov.Model):
 
     fuse_cache_reorder(ov_model, not_kv_inputs, key_value_input_names, batch_dim)
 
+<<<<<<< HEAD
     normalized_config = NormalizedConfigManager.get_normalized_config_class(config.model_type)(config)
     num_attention_heads = normalized_config.num_attention_heads if config.model_type == "bloom" else 1
+=======
+    normalized_config = NormalizedConfigManager.get_normalized_config_class(model.config.model_type)(model.config)
+    num_attention_heads = normalized_config.num_attention_heads if model.config.model_type == 'bloom' else 1
+>>>>>>> Fix for bloom family
 
     make_stateful(
         ov_model, not_kv_inputs, key_value_input_names, key_value_output_names, batch_dim, num_attention_heads, None