fix

Author: echarlaix

optimum/intel/ipex/inference.py

@@ -31,10 +31,24 @@
     IPEXModelForMaskedLM,
     IPEXModelForSequenceClassification,
     IPEXModelForTokenClassification,
+    IPEXBloomForCausalLM,
+    IPEXMPTForCausalLM,
+    IPEXOPTForCausalLM,
+    IPEXGPTBigCodeForCausalLM,
 )
+
+
 from .utils import _HEAD_TO_AUTOMODELS
 
 
+_MODEL_TYPE_TO_AUTOMODELS = {
+    "bloom": IPEXBloomForCausalLM,
+    "mpt": IPEXMPTForCausalLM,
+    "opt": IPEXOPTForCausalLM,
+    "big_code": IPEXGPTBigCodeForCausalLM,
+}
+
+
 logger = logging.getLogger(__name__)
 
 IPEX_NOT_AVAILABLE_ERROR_MSG = (
@@ -131,7 +145,14 @@ def __enter__(self):
                         )
                         if task in _HEAD_TO_AUTOMODELS:
                             model = jit_trace(model, task, use_cache)
-                            model = eval(_HEAD_TO_AUTOMODELS[task])(model, self._original.config, use_cache=use_cache)
+                            model_type = getattr(self._original.config, "model_type", "").replace("_", "-")
+
+                            if task == "text-generation" and model_type in _MODEL_TYPE_TO_AUTOMODELS.keys():
+                                auto_model_class = _MODEL_TYPE_TO_AUTOMODELS[task]
+                            else:
+                                auto_model_class = eval(_HEAD_TO_AUTOMODELS[task])
+
+                            model = auto_model_class(model, self._original.config, use_cache=use_cache)
 
                     # Enable automatic mixed precision (AMP) if we are going to target `bfloat16`
                     with torch.cpu.amp.autocast(enabled=self._dtype == torch.bfloat16):

optimum/intel/ipex/modeling_base.py

@@ -158,17 +158,10 @@ def _from_pretrained(
 
         model = torch.jit.load(model_cache_path)
         torch.jit.freeze(model.eval())
-
+        model_type = config.model_type.replace("_", "-")
         init_cls = cls
-        if cls is IPEXModelForCausalLM:
-            if config.model_type == "bloom":
-                init_cls = IPEXBloomForCausalLM
-            elif config.model_type == "mpt":
-                init_cls = IPEXMPTForCausalLM
-            elif config.model_type == "opt":
-                init_cls = IPEXOPTForCausalLM
-            elif config.model_type == "gpt_bigcode":
-                init_cls = IPEXGPTBigCodeForCausalLM
+        if cls.export_feature == "text-generation" and model_type in _MODEL_TYPE_TO_AUTOMODELS:
+            init_cls = _MODEL_TYPE_TO_AUTOMODELS[model_type]
 
         return init_cls(model, config=config, model_save_dir=model_save_dir, **kwargs)
 
@@ -535,3 +528,11 @@ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwarg
             "position_ids": None,
             "attention_mask": attention_mask,
         }
+
+
+_MODEL_TYPE_TO_AUTOMODELS = {
+    "bloom": IPEXBloomForCausalLM,
+    "mpt": IPEXMPTForCausalLM,
+    "opt": IPEXOPTForCausalLM,
+    "big-code": IPEXGPTBigCodeForCausalLM,
+}

optimum/intel/ipex/utils.py

@@ -17,5 +17,5 @@
     "text-generation": "IPEXModelForCausalLM",
     "text-classification": "IPEXModelForSequenceClassification",
     "token-classification": "IPEXModelForTokenClassification",
-    "question-answering": "IPEXModelForQuestionAnswering",
+    # "question-answering": "IPEXModelForQuestionAnswering",
 }

tests/ipex/test_inference.py

@@ -28,7 +28,7 @@
 )
 
 from optimum.intel import inference_mode as ipex_inference_mode
-from optimum.intel.generation.modeling import TSModelForCausalLM
+from optimum.intel.ipex.modeling_base import IPEXModel
 
 
 MODEL_NAMES = {
@@ -112,6 +112,6 @@ def test_text_generation_pipeline_inference(self, model_arch):
             text_generator, dtype=model.config.torch_dtype, verbose=False, jit=True
         ) as ipex_text_generator:
             output_ipex = ipex_text_generator(inputs)
-        self.assertTrue(isinstance(ipex_text_generator.model._optimized, TSModelForCausalLM))
+        self.assertTrue(isinstance(ipex_text_generator.model._optimized, IPEXModel))
         self.assertTrue(isinstance(ipex_text_generator.model._optimized.model, torch.jit.RecursiveScriptModule))
         self.assertEqual(output[0]["generated_text"], output_ipex[0]["generated_text"])