From caf2760bdf678c692c089316c94b16cfe5b8cffa Mon Sep 17 00:00:00 2001
From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com>
Date: Sat, 9 Mar 2024 17:20:32 +0000
Subject: [PATCH] support phi

---
 optimum/exporters/neuron/model_configs.py  | 13 +++++++++++++
 optimum/exporters/neuron/model_wrappers.py | 13 +++++++++++++
 tests/exporters/exporters_utils.py         |  1 +
 tests/inference/inference_utils.py         |  1 +
 4 files changed, 28 insertions(+)

diff --git a/optimum/exporters/neuron/model_configs.py b/optimum/exporters/neuron/model_configs.py
index 3d0808c1c..d6991b89d 100644
--- a/optimum/exporters/neuron/model_configs.py
+++ b/optimum/exporters/neuron/model_configs.py
@@ -43,6 +43,7 @@
     VisionNeuronConfig,
 )
 from .model_wrappers import (
+    NoCacheModelWrapper,
     SentenceTransformersCLIPNeuronWrapper,
     SentenceTransformersTransformerNeuronWrapper,
     T5DecoderWrapper,
@@ -122,6 +123,18 @@ class MobileBertNeuronConfig(BertNeuronConfig):
     pass
 
 
+@register_in_tasks_manager("phi", *["feature-extraction", "text-classification", "token-classification"])
+class PhiNeuronConfig(ElectraNeuronConfig):
+    CUSTOM_MODEL_WRAPPER = NoCacheModelWrapper
+
+    @property
+    def inputs(self) -> List[str]:
+        return ["input_ids", "attention_mask"]
+
+    def patch_model_for_export(self, model, dummy_inputs):
+        return self.CUSTOM_MODEL_WRAPPER(model, list(dummy_inputs.keys()))
+
+
 @register_in_tasks_manager("roformer", *COMMON_TEXT_TASKS)
 class RoFormerNeuronConfig(ElectraNeuronConfig):
     pass
diff --git a/optimum/exporters/neuron/model_wrappers.py b/optimum/exporters/neuron/model_wrappers.py
index 711ff509d..261701c76 100644
--- a/optimum/exporters/neuron/model_wrappers.py
+++ b/optimum/exporters/neuron/model_wrappers.py
@@ -374,3 +374,16 @@ def forward(self, input_ids, pixel_values, attention_mask):
             text_embeds = self.model[1:](text_embeds)
 
         return (text_embeds, image_embeds)
+
+
+class NoCacheModelWrapper(torch.nn.Module):
+    def __init__(self, model: "PreTrainedModel", input_names: List[str]):
+        super().__init__()
+        self.model = model
+        self.input_names = input_names
+
+    def forward(self, *input):
+        ordered_inputs = dict(zip(self.input_names, input))
+        outputs = self.model(use_cache=False, **ordered_inputs)
+
+        return outputs
diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
index e1b5e10f2..483c6f856 100644
--- a/tests/exporters/exporters_utils.py
+++ b/tests/exporters/exporters_utils.py
@@ -27,6 +27,7 @@
     "flaubert": "flaubert/flaubert_small_cased",
     "mobilebert": "hf-internal-testing/tiny-random-MobileBertModel",
     "mpnet": "hf-internal-testing/tiny-random-MPNetModel",
+    "phi": "hf-internal-testing/tiny-random-PhiModel",
     "roberta": "hf-internal-testing/tiny-random-RobertaModel",
     "roformer": "hf-internal-testing/tiny-random-RoFormerModel",
     "xlm": "hf-internal-testing/tiny-random-XLMModel",
diff --git a/tests/inference/inference_utils.py b/tests/inference/inference_utils.py
index f64406a6a..c5e63fd01 100644
--- a/tests/inference/inference_utils.py
+++ b/tests/inference/inference_utils.py
@@ -42,6 +42,7 @@
     "latent-consistency": "echarlaix/tiny-random-latent-consistency",
     "mobilebert": "hf-internal-testing/tiny-random-MobileBertModel",
     "mpnet": "hf-internal-testing/tiny-random-MPNetModel",
+    "phi": "hf-internal-testing/tiny-random-PhiModel",
     "roberta": "hf-internal-testing/tiny-random-RobertaModel",
     "roformer": "hf-internal-testing/tiny-random-RoFormerModel",
     "stable-diffusion": "hf-internal-testing/tiny-stable-diffusion-torch",