Skip to content

Commit 2b8c6b1

Browse files
committed
add export config for starcoder2
1 parent d2f9fdb commit 2b8c6b1

File tree

3 files changed

+13
-0
lines changed

3 files changed

+13
-0
lines changed

optimum/exporters/openvino/model_configs.py

+11
Original file line numberDiff line numberDiff line change
@@ -389,3 +389,14 @@ def patch_model_for_export(
389389
self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
390390
) -> "ModelPatcher":
391391
return QwenModelPatcher(self, model, model_kwargs=model_kwargs)
392+
393+
394+
@register_in_tasks_manager(
395+
"starcoder2", *["text-generation", "text-generation-with-past"], library_name="transformers"
396+
)
397+
class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig):
398+
DEFAULT_ONNX_OPSET = 14
399+
400+
DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator)
401+
DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator
402+
NORMALIZED_CONFIG_CLASS = NormalizedTextConfig

tests/openvino/test_modeling.py

+1
Original file line numberDiff line numberDiff line change
@@ -501,6 +501,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
501501
"qwen",
502502
"qwen2",
503503
"stablelm",
504+
# "starcoder2", # TODO: enable with next transformers release
504505
)
505506
GENERATION_LENGTH = 100
506507
IS_SUPPORT_STATEFUL = is_openvino_version(">=", "2023.3")

tests/openvino/utils_tests.py

+1
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@
8484
"stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
8585
"stable-diffusion-xl-refiner": "echarlaix/tiny-random-stable-diffusion-xl-refiner",
8686
"stablelm": "hf-internal-testing/tiny-random-StableLmForCausalLM",
87+
"starcoder2": "hf-internal-testing/tiny-random-Starcoder2ForCausalLM",
8788
"latent-consistency": "echarlaix/tiny-random-latent-consistency",
8889
"sew": "hf-internal-testing/tiny-random-SEWModel",
8990
"sew_d": "asapp/sew-d-tiny-100k-ft-ls100h",

0 commit comments

Comments
 (0)