Merge pull request #619 from eaidova/ea/starcoder2

add export config for starcoder2
huggingface · Mar 20, 2024 · 086fae3 · 086fae3
2 parents d2f9fdb + 2b8c6b1
commit 086fae3
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py
@@ -389,3 +389,14 @@ def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> "ModelPatcher":
         return QwenModelPatcher(self, model, model_kwargs=model_kwargs)
+
+
+@register_in_tasks_manager(
+    "starcoder2", *["text-generation", "text-generation-with-past"], library_name="transformers"
+)
+class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig):
+    DEFAULT_ONNX_OPSET = 14
+
+    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator)
+    DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator
+    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
diff --git a/tests/openvino/test_modeling.py b/tests/openvino/test_modeling.py
@@ -501,6 +501,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "qwen",
         "qwen2",
         "stablelm",
+        # "starcoder2", # TODO: enable with next transformers release
     )
     GENERATION_LENGTH = 100
     IS_SUPPORT_STATEFUL = is_openvino_version(">=", "2023.3")

diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py
@@ -84,6 +84,7 @@
     "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
     "stable-diffusion-xl-refiner": "echarlaix/tiny-random-stable-diffusion-xl-refiner",
     "stablelm": "hf-internal-testing/tiny-random-StableLmForCausalLM",
+    "starcoder2": "hf-internal-testing/tiny-random-Starcoder2ForCausalLM",
     "latent-consistency": "echarlaix/tiny-random-latent-consistency",
     "sew": "hf-internal-testing/tiny-random-SEWModel",
     "sew_d": "asapp/sew-d-tiny-100k-ft-ls100h",