support phi

huggingface · Mar 9, 2024 · caf2760 · caf2760
1 parent 8f84127
commit caf2760
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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",