Add IPEX model for question answering (#534)

* add IPEX model for QA task

* add fix

optimum/intel/__init__.py

@@ -47,6 +47,7 @@
         "IPEXModelForSequenceClassification",
         "IPEXModelForMaskedLM",
         "IPEXModelForTokenClassification",
+        "IPEXModelForQuestionAnswering",
     ]
 
 
@@ -160,6 +161,7 @@
         from .ipex import (
             IPEXModelForCausalLM,
             IPEXModelForMaskedLM,
+            IPEXModelForQuestionAnswering,
             IPEXModelForSequenceClassification,
             IPEXModelForTokenClassification,
             inference_mode,

optimum/intel/ipex/__init__.py

@@ -1,6 +1,7 @@
 from optimum.intel.ipex.modeling_base import (
     IPEXModelForCausalLM,
     IPEXModelForMaskedLM,
+    IPEXModelForQuestionAnswering,
     IPEXModelForSequenceClassification,
     IPEXModelForTokenClassification,
 )

optimum/intel/ipex/inference.py

@@ -35,6 +35,7 @@
     IPEXMPTForCausalLM,
     IPEXOPTForCausalLM,
     IPEXGPTBigCodeForCausalLM,
+    IPEXModelForQuestionAnswering,
 )
 
 

optimum/intel/ipex/modeling_base.py

@@ -27,6 +27,7 @@
     AutoModel,
     AutoModelForCausalLM,
     AutoModelForMaskedLM,
+    AutoModelForQuestionAnswering,
     AutoModelForSequenceClassification,
     AutoModelForTokenClassification,
     GenerationConfig,
@@ -171,7 +172,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
 
     def forward(self, *args, **kwargs):
         outputs = self.model(*args, **kwargs)
-        return ModelOutput(logits=outputs["logits"] if isinstance(outputs, dict) else outputs[0])
+        return ModelOutput(**outputs) if isinstance(outputs, dict) else ModelOutput(logits=outputs[0])
 
     def eval(self):
         self.model.eval()
@@ -205,6 +206,17 @@ class IPEXModelForTokenClassification(IPEXModel):
     export_feature = "token-classification"
 
 
+class IPEXModelForQuestionAnswering(IPEXModel):
+    auto_model_class = AutoModelForQuestionAnswering
+    export_feature = "question-answering"
+
+    def forward(self, *args, **kwargs):
+        outputs = self.model(*args, **kwargs)
+        start_logits = outputs["start_logits"] if isinstance(outputs, dict) else outputs[0]
+        end_logits = outputs["end_logits"] if isinstance(outputs, dict) else outputs[1]
+        return ModelOutput(start_logits=start_logits, end_logits=end_logits)
+
+
 class IPEXModelForCausalLM(IPEXModel, GenerationMixin):
     auto_model_class = AutoModelForCausalLM
     export_feature = "text-generation"

optimum/intel/ipex/utils.py

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

optimum/intel/utils/dummy_ipex_objects.py

@@ -64,3 +64,14 @@ def __init__(self, *args, **kwargs):
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["ipex"])
+
+
+class IPEXModelForQuestionAnswering(metaclass=DummyObject):
+    _backends = ["ipex"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["ipex"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["ipex"])

tests/ipex/test_inference.py

@@ -78,7 +78,7 @@ def test_question_answering_pipeline_inference(self, model_arch):
             outputs_ipex = ipex_pipe(
                 question="Where was HuggingFace founded ?", context="HuggingFace was founded in Paris."
             )
-        # self.assertTrue(isinstance(ipex_pipe.model._optimized.model, torch.jit.RecursiveScriptModule))
+        self.assertTrue(isinstance(ipex_pipe.model._optimized.model, torch.jit.RecursiveScriptModule))
         self.assertEqual(outputs["start"], outputs_ipex["start"])
         self.assertEqual(outputs["end"], outputs_ipex["end"])
 

tests/ipex/test_modeling.py

@@ -20,6 +20,7 @@
 from parameterized import parameterized
 from transformers import (
     AutoModelForCausalLM,
+    AutoModelForQuestionAnswering,
     AutoModelForSequenceClassification,
     AutoTokenizer,
     PretrainedConfig,
@@ -28,7 +29,7 @@
 )
 
 from optimum.exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS
-from optimum.intel import IPEXModelForCausalLM, IPEXModelForSequenceClassification
+from optimum.intel import IPEXModelForCausalLM, IPEXModelForQuestionAnswering, IPEXModelForSequenceClassification
 
 
 SEED = 42
@@ -118,6 +119,46 @@ def test_pipeline(self, model_arch):
         self.assertIsInstance(outputs[0]["label"], str)
 
 
+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 IPEXModelForCausalLMTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES = (
         "bart",