diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index d8fa4ef6ec..2712f9da31 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -66,6 +66,7 @@ NormalizedTimeSeriesForecastingConfig, NormalizedVisionConfig, PerceiverDummyInputGenerator, + TimesFMDummyInputGenerator, VitPoseDummyInputGenerator, is_diffusers_available, is_diffusers_version, @@ -2641,6 +2642,25 @@ class EncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): DEFAULT_ONNX_OPSET = 14 # uses SDPA in Transformers, hence opset>=14. +class TimesFMOnnxConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTimeSeriesForecastingConfig + MIN_TRANSFORMERS_VERSION = version.parse("4.52.0") + DUMMY_INPUT_GENERATOR_CLASSES = (TimesFMDummyInputGenerator,) + DEFAULT_ONNX_OPSET = 14 # uses SDPA in Transformers, needs opset>=14 + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return {"past_values": {0: "batch_size", 1: "sequence_length"}} + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "last_hidden_state": {0: "batch_size"}, + "mean_predictions": {0: "batch_size"}, + "full_predictions": {0: "batch_size"}, + } + + class PatchTSTOnnxConfig(OnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTimeSeriesForecastingConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyPatchTSTInputGenerator,) diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index 21f05135d8..bbe83bd6d5 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -329,6 +329,7 @@ class TasksManager: ("pt", "visual-bert", "question-answering"): ("transformers", "VisualBertForQuestionAnswering"), # VisionEncoderDecoderModel is not registered in AutoModelForDocumentQuestionAnswering ("pt", "vision-encoder-decoder", "document-question-answering"): ("transformers", "VisionEncoderDecoderModel"), + ("pt", "timesfm", "time-series-forecasting"): ("transformers", "TimesFmModelForPrediction"), ("pt", "vitpose", "keypoint-detection"): ("transformers", "VitPoseForPoseEstimation"), } @@ -996,6 +997,10 @@ class TasksManager: "text-classification", onnx="Qwen2OnnxConfig", ), + "timesfm": supported_tasks_mapping( + "time-series-forecasting", + onnx="TimesFMOnnxConfig", + ), "llama": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", diff --git a/optimum/utils/__init__.py b/optimum/utils/__init__.py index f945ac14f3..0dcf4c93bc 100644 --- a/optimum/utils/__init__.py +++ b/optimum/utils/__init__.py @@ -92,6 +92,7 @@ MistralDummyPastKeyValuesGenerator, MultiQueryPastKeyValuesGenerator, PerceiverDummyInputGenerator, + TimesFMDummyInputGenerator, VitPoseDummyInputGenerator, ) from .modeling_utils import recurse_getattr, recurse_setattr diff --git a/optimum/utils/input_generators.py b/optimum/utils/input_generators.py index 40cf2cde49..2fae73288e 100644 --- a/optimum/utils/input_generators.py +++ b/optimum/utils/input_generators.py @@ -1658,3 +1658,28 @@ class PerceiverDummyInputGenerator(DummyVisionStaticInputGenerator): class VitPoseDummyInputGenerator(DummyVisionStaticInputGenerator): pass + + +class TimesFMDummyInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ("past_values",) + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + **kwargs, + ): + self.task = task + self.normalized_config = normalized_config + self.batch_size = batch_size + self.context_length = normalized_config.context_length + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + return self.random_float_tensor( + shape=[self.batch_size, self.context_length], + min_value=-1, + max_value=1, + framework=framework, + dtype=float_dtype, + ) diff --git a/tests/exporters/utils.py b/tests/exporters/utils.py index d2310a1d2d..fdaa702ef0 100644 --- a/tests/exporters/utils.py +++ b/tests/exporters/utils.py @@ -171,6 +171,7 @@ "swin2sr": "hf-internal-testing/tiny-random-Swin2SRModel", "t5": "hf-internal-testing/tiny-random-t5", "table-transformer": "hf-internal-testing/tiny-random-TableTransformerModel", + "timesfm": "huggingface/timesfm-tourism-monthly", "vit": "hf-internal-testing/tiny-random-vit", "vit-mae": "hf-internal-testing/tiny-random-ViTMAEModel", "vit-msn": "hf-internal-testing/tiny-random-ViTMSNForImageClassification",