chronos-forecasting 1.3.0__py3-none-any.whl

chronos/__init__.py

@@ -0,0 +1,24 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .base import BaseChronosPipeline, ForecastType
+from .chronos import (
+    ChronosConfig,
+    ChronosModel,
+    ChronosPipeline,
+    ChronosTokenizer,
+    MeanScaleUniformBins,
+)
+from .chronos_bolt import ChronosBoltConfig, ChronosBoltPipeline
+
+__all__ = [
+    "BaseChronosPipeline",
+    "ForecastType",
+    "ChronosConfig",
+    "ChronosModel",
+    "ChronosPipeline",
+    "ChronosTokenizer",
+    "MeanScaleUniformBins",
+    "ChronosBoltConfig",
+    "ChronosBoltPipeline",
+]

chronos/base.py

@@ -0,0 +1,162 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Authors: Caner Turkmen <atturkm@amazon.com>, Abdul Fatir Ansari <ansarnd@amazon.com>, Lorenzo Stella <stellalo@amazon.com>
+# Original source:
+# https://github.com/autogluon/autogluon/blob/f57beb26cb769c6e0d484a6af2b89eab8aee73a8/timeseries/src/autogluon/timeseries/models/chronos/pipeline/base.py
+
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+import torch
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+from .utils import left_pad_and_stack_1D
+
+
+class ForecastType(Enum):
+    SAMPLES = "samples"
+    QUANTILES = "quantiles"
+
+
+class PipelineRegistry(type):
+    REGISTRY: Dict[str, "PipelineRegistry"] = {}
+
+    def __new__(cls, name, bases, attrs):
+        """See, https://github.com/faif/python-patterns."""
+        new_cls = type.__new__(cls, name, bases, attrs)
+        if name is not None:
+            cls.REGISTRY[name] = new_cls
+
+        return new_cls
+
+
+class BaseChronosPipeline(metaclass=PipelineRegistry):
+    forecast_type: ForecastType
+    dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32}
+
+    def __init__(self, inner_model: "PreTrainedModel"):
+        """
+        Parameters
+        ----------
+        inner_model : PreTrainedModel
+            A hugging-face transformers PreTrainedModel, e.g., T5ForConditionalGeneration
+        """
+        # for easy access to the inner HF-style model
+        self.inner_model = inner_model
+
+    def _prepare_and_validate_context(
+        self, context: Union[torch.Tensor, List[torch.Tensor]]
+    ):
+        if isinstance(context, list):
+            context = left_pad_and_stack_1D(context)
+        assert isinstance(context, torch.Tensor)
+        if context.ndim == 1:
+            context = context.unsqueeze(0)
+        assert context.ndim == 2
+
+        return context
+
+    def predict(
+        self,
+        context: Union[torch.Tensor, List[torch.Tensor]],
+        prediction_length: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Get forecasts for the given time series.
+
+        Parameters
+        ----------
+        context
+            Input series. This is either a 1D tensor, or a list
+            of 1D tensors, or a 2D tensor whose first dimension
+            is batch. In the latter case, use left-padding with
+            ``torch.nan`` to align series of different lengths.
+        prediction_length
+            Time steps to predict. Defaults to a model-dependent
+            value if not given.
+
+        Returns
+        -------
+        forecasts
+            Tensor containing forecasts. The layout and meaning
+            of the forecasts values depends on ``self.forecast_type``.
+        """
+        raise NotImplementedError()
+
+    def predict_quantiles(
+        self,
+        context: Union[torch.Tensor, List[torch.Tensor]],
+        prediction_length: Optional[int] = None,
+        quantile_levels: List[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Get quantile and mean forecasts for given time series.
+
+        Parameters
+        ----------
+        context : Union[torch.Tensor, List[torch.Tensor]]
+            Input series. This is either a 1D tensor, or a list
+            of 1D tensors, or a 2D tensor whose first dimension
+            is batch. In the latter case, use left-padding with
+            ``torch.nan`` to align series of different lengths.
+        prediction_length : Optional[int], optional
+            Time steps to predict. Defaults to a model-dependent
+            value if not given.
+        quantile_levels : List[float], optional
+            Quantile levels to compute, by default [0.1, 0.2, ..., 0.9]
+
+        Returns
+        -------
+        quantiles
+            Tensor containing quantile forecasts. Shape
+            (batch_size, prediction_length, num_quantiles)
+        mean
+            Tensor containing mean (point) forecasts. Shape
+            (batch_size, prediction_length)
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+        *model_args,
+        **kwargs,
+    ):
+        """
+        Load the model, either from a local path or from the HuggingFace Hub.
+        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
+        from ``transformers``.
+        """
+        from transformers import AutoConfig
+
+        torch_dtype = kwargs.get("torch_dtype", "auto")
+        if torch_dtype != "auto" and isinstance(torch_dtype, str):
+            kwargs["torch_dtype"] = cls.dtypes[torch_dtype]
+
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        is_valid_config = hasattr(config, "chronos_pipeline_class") or hasattr(
+            config, "chronos_config"
+        )
+
+        if not is_valid_config:
+            raise ValueError("Not a Chronos config file")
+
+        pipeline_class_name = getattr(
+            config, "chronos_pipeline_class", "ChronosPipeline"
+        )
+        class_ = PipelineRegistry.REGISTRY.get(pipeline_class_name)
+        if class_ is None:
+            raise ValueError(
+                f"Trying to load unknown pipeline class: {pipeline_class_name}"
+            )
+
+        return class_.from_pretrained(  # type: ignore[attr-defined]
+            pretrained_model_name_or_path, *model_args, **kwargs
+        )