Fix TS model for BLOOM architecture (#344)

Author: echarlaix

optimum/intel/generation/modeling.py

@@ -31,6 +31,7 @@
 
 from ..utils.constant import _TASK_ALIASES
 from ..utils.import_utils import is_torch_version, is_transformers_version
+from ..utils.modeling_utils import _prepare_attn_mask, _prepare_decoder_attention_mask
 
 
 if is_transformers_version("<", "4.25.0"):
@@ -266,6 +267,13 @@ def _from_transformers(
         }
 
         model = TasksManager.get_model_from_task(task, model_id, **model_kwargs)
+
+        if model.config.model_type == "bloom":
+            model.transformer._prepare_attn_mask = _prepare_attn_mask
+
+        if model.config.model_type == "llama":
+            model.model._prepare_decoder_attention_mask = _prepare_decoder_attention_mask
+
         traced_model = jit_trace(model, task, use_cache)
         save_dir = TemporaryDirectory()
         save_dir_path = Path(save_dir.name)

optimum/intel/openvino/modeling_decoder.py

@@ -30,8 +30,8 @@
 from optimum.utils import NormalizedConfigManager
 
 from ..utils.import_utils import is_transformers_version
+from ..utils.modeling_utils import _prepare_attn_mask, _prepare_decoder_attention_mask
 from .modeling import _TOKENIZER_FOR_DOC, INPUTS_DOCSTRING, MODEL_START_DOCSTRING, OVModel
-from .modeling_utils import _prepare_attn_mask, _prepare_decoder_attention_mask
 from .utils import ONNX_WEIGHTS_NAME
 
 

optimum/intel/openvino/modeling_utils.py optimum/intel/utils/modeling_utils.py

@@ -66,8 +66,7 @@ def test_compare_to_transformers(self, model_arch):
         with torch.no_grad():
             trfs_outputs = trfs_model(**tokens)
         # Compare outputs with original transformers model
-        atol = 1e-1 if model_arch == "bloom" else 1e-4
-        self.assertTrue(torch.allclose(outputs.logits, trfs_outputs.logits, atol=atol))
+        self.assertTrue(torch.allclose(outputs.logits, trfs_outputs.logits, atol=1e-4))
         # Compare outputs with loaded model
         with tempfile.TemporaryDirectory() as tmpdirname:
             model.save_pretrained(tmpdirname)