IPEX decoder model fix (#539)

Author: echarlaix

optimum/intel/ipex/inference.py

@@ -36,20 +36,13 @@
     IPEXOPTForCausalLM,
     IPEXGPTBigCodeForCausalLM,
     IPEXModelForQuestionAnswering,
+    _MODEL_TYPE_TO_AUTOMODELS,
 )
 
 
 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 = (
@@ -149,7 +142,7 @@ def __enter__(self):
                             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]
+                                auto_model_class = _MODEL_TYPE_TO_AUTOMODELS[model_type]
                             else:
                                 auto_model_class = eval(_HEAD_TO_AUTOMODELS[task])
 

optimum/intel/ipex/modeling_base.py

@@ -600,5 +600,5 @@ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwarg
     "bloom": IPEXBloomForCausalLM,
     "mpt": IPEXMPTForCausalLM,
     "opt": IPEXOPTForCausalLM,
-    "big-code": IPEXGPTBigCodeForCausalLM,
+    "gpt-bigcode": IPEXGPTBigCodeForCausalLM,
 }

tests/ipex/test_inference.py

@@ -42,6 +42,8 @@
     "gpt_neox": "hf-internal-testing/tiny-random-GPTNeoXForCausalLM",
     "gpt_bigcode": "hf-internal-testing/tiny-random-GPTBigCodeModel",
     "llama": "fxmarty/tiny-llama-fast-tokenizer",
+    "opt": "hf-internal-testing/tiny-random-OPTModel",
+    "mpt": "hf-internal-testing/tiny-random-MptForCausalLM",
 }
 
 _CLASSIFICATION_TASK_TO_AUTOMODELS = {
@@ -57,7 +59,7 @@ class IPEXIntegrationTest(unittest.TestCase):
         "roberta",
     )
 
-    TEXT_GENERATION_SUPPORTED_ARCHITECTURES = ("gptj", "gpt2", "gpt_neo", "gpt_bigcode", "llama")
+    TEXT_GENERATION_SUPPORTED_ARCHITECTURES = ("gptj", "gpt2", "gpt_neo", "gpt_bigcode", "llama", "opt", "mpt")
 
     QA_SUPPORTED_ARCHITECTURES = (
         "bert",