Add task

Author: echarlaix

optimum/intel/__init__.py

@@ -48,10 +48,11 @@
         "IPEXModelForMaskedLM",
         "IPEXModelForTokenClassification",
         "IPEXModelForQuestionAnswering",
+        "IPEXModelForImageClassification",
+        "IPEXModelForAudioClassification",
         "IPEXModel",
     ]
 
-
 try:
     if not (is_openvino_available() and is_nncf_available()):
         raise OptionalDependencyNotAvailable()
@@ -161,7 +162,9 @@
     else:
         from .ipex import (
             IPEXModel,
+            IPEXModelForAudioClassification,
             IPEXModelForCausalLM,
+            IPEXModelForImageClassification,
             IPEXModelForMaskedLM,
             IPEXModelForQuestionAnswering,
             IPEXModelForSequenceClassification,

optimum/intel/generation/modeling.py

@@ -66,13 +66,12 @@ def prepare_jit_inputs(model: PreTrainedModel, task: str, use_cache: bool = Fals
 
 def jit_trace(model: PreTrainedModel, task: str, use_cache: bool = False):
     model_inputs = prepare_jit_inputs(model, task, use_cache)
-    model.config.return_dict = False
+    model.config.return_dict = "text-generation" not in task
+
     # check if the model_inputs is correct.
     model(**model_inputs)
 
     torch._C._jit_set_texpr_fuser_enabled(False)
-    if "past_key_values" in model_inputs.keys():
-        model.config.return_dict = False
     if is_torch_version(">=", "2.1.0"):
         traced_model = torch.jit.trace(model, example_kwarg_inputs=model_inputs, strict=False)
     else:

optimum/intel/ipex/__init__.py

@@ -1,6 +1,8 @@
 from optimum.intel.ipex.modeling_base import (
     IPEXModel,
+    IPEXModelForAudioClassification,
     IPEXModelForCausalLM,
+    IPEXModelForImageClassification,
     IPEXModelForMaskedLM,
     IPEXModelForQuestionAnswering,
     IPEXModelForSequenceClassification,

optimum/intel/ipex/modeling_base.py

@@ -25,7 +25,9 @@
 from transformers import (
     AutoConfig,
     AutoModel,
+    AutoModelForAudioClassification,
     AutoModelForCausalLM,
+    AutoModelForImageClassification,
     AutoModelForMaskedLM,
     AutoModelForQuestionAnswering,
     AutoModelForSequenceClassification,
@@ -196,14 +198,24 @@ class IPEXModelForSequenceClassification(IPEXModel):
     export_feature = "text-classification"
 
 
+class IPEXModelForTokenClassification(IPEXModel):
+    auto_model_class = AutoModelForTokenClassification
+    export_feature = "token-classification"
+
+
 class IPEXModelForMaskedLM(IPEXModel):
     auto_model_class = AutoModelForMaskedLM
     export_feature = "fill-mask"
 
 
-class IPEXModelForTokenClassification(IPEXModel):
-    auto_model_class = AutoModelForTokenClassification
-    export_feature = "token-classification"
+class IPEXModelForImageClassification(IPEXModel):
+    auto_model_class = AutoModelForImageClassification
+    export_feature = "image-classification"
+
+
+class IPEXModelForAudioClassification(IPEXModel):
+    auto_model_class = AutoModelForAudioClassification
+    export_feature = "audio-classification"
 
 
 class IPEXModelForQuestionAnswering(IPEXModel):

optimum/intel/utils/dummy_ipex_objects.py

@@ -22,6 +22,17 @@ def __init__(self, *args, **kwargs):
         requires_backends(self, ["ipex"])
 
 
+class IPEXModel(metaclass=DummyObject):
+    _backends = ["ipex"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["ipex"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["ipex"])
+
+
 class IPEXModelForSequenceClassification(metaclass=DummyObject):
     _backends = ["ipex"]
 
@@ -77,7 +88,29 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["ipex"])
 
 
-class IPEXModel(metaclass=DummyObject):
+class IPEXModelForMaskedLM(metaclass=DummyObject):
+    _backends = ["ipex"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["ipex"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["ipex"])
+
+
+class IPEXModelForImageClassification(metaclass=DummyObject):
+    _backends = ["ipex"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["ipex"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["ipex"])
+
+
+class IPEXModelForAudioClassification(metaclass=DummyObject):
     _backends = ["ipex"]
 
     def __init__(self, *args, **kwargs):

tests/ipex/test_modeling.py

@@ -30,7 +30,10 @@
 from optimum.exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS
 from optimum.intel import (
     IPEXModel,
+    IPEXModelForAudioClassification,
     IPEXModelForCausalLM,
+    IPEXModelForImageClassification,
+    IPEXModelForMaskedLM,
     IPEXModelForQuestionAnswering,
     IPEXModelForSequenceClassification,
     IPEXModelForTokenClassification,
@@ -110,7 +113,9 @@ def test_compare_to_transformers(self, model_arch):
             transformers_outputs = transformers_model(**tokens)
         outputs = ipex_model(**tokens)
         # Compare tensor outputs
-        self.assertTrue(torch.allclose(outputs.logits, transformers_outputs.logits, atol=1e-4))
+        for output_name in {"logits", "last_hidden_state"}:
+            if output_name in transformers_outputs:
+                self.assertTrue(torch.allclose(outputs[output_name], transformers_outputs[output_name], atol=1e-4))
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     def test_pipeline(self, model_arch):
@@ -119,7 +124,7 @@ def test_pipeline(self, model_arch):
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         pipe = pipeline(self.IPEX_MODEL_CLASS.export_feature, model=model, tokenizer=tokenizer)
         text = "This restaurant is awesome"
-        outputs = pipe(text)
+        _ = pipe(text)
 
         self.assertEqual(pipe.device, model.device)
 
@@ -132,44 +137,8 @@ class IPEXModelForTokenClassificationTest(IPEXModelTest):
     IPEX_MODEL_CLASS = IPEXModelForSequenceClassification
 
 
-class IPEXModelForQuestionAnsweringTest(unittest.TestCase):
-    SUPPORTED_ARCHITECTURES = (
-        "bert",
-        "distilbert",
-        "roberta",
-    )
-
-    @parameterized.expand(SUPPORTED_ARCHITECTURES)
-    def test_compare_to_transformers(self, model_arch):
-        model_id = MODEL_NAMES[model_arch]
-        set_seed(SEED)
-        ipex_model = IPEXModelForQuestionAnswering.from_pretrained(model_id, export=True)
-        self.assertIsInstance(ipex_model.config, PretrainedConfig)
-        transformers_model = AutoModelForQuestionAnswering.from_pretrained(model_id)
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        inputs = "This is a sample input"
-        tokens = tokenizer(inputs, return_tensors="pt")
-        with torch.no_grad():
-            transformers_outputs = transformers_model(**tokens)
-        outputs = ipex_model(**tokens)
-        self.assertIn("start_logits", outputs)
-        self.assertIn("end_logits", outputs)
-        # Compare tensor outputs
-        self.assertTrue(torch.allclose(outputs.start_logits, transformers_outputs.start_logits, atol=1e-4))
-        self.assertTrue(torch.allclose(outputs.end_logits, transformers_outputs.end_logits, atol=1e-4))
-
-    @parameterized.expand(SUPPORTED_ARCHITECTURES)
-    def test_pipeline(self, model_arch):
-        model_id = MODEL_NAMES[model_arch]
-        model = IPEXModelForQuestionAnswering.from_pretrained(model_id, export=True)
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        pipe = pipeline("question-answering", model=model, tokenizer=tokenizer)
-        question = "What's my name?"
-        context = "My Name is Sasha and I live in Lyon."
-        outputs = pipe(question, context)
-        self.assertEqual(pipe.device, model.device)
-        self.assertGreaterEqual(outputs["score"], 0.0)
-        self.assertIsInstance(outputs["answer"], str)
+class IPEXModelForMaskedLMTest(IPEXModelTest):
+    IPEX_MODEL_CLASS = IPEXModelForMaskedLM
 
 
 class IPEXModelForQuestionAnsweringTest(unittest.TestCase):