autogluon.timeseries 1.1.2b20241109__py3-none-any.whl → 1.1.2b20241112__py3-none-any.whl

autogluon/timeseries/models/chronos/{pipeline.py → pipeline/chronos.py}

@@ -11,14 +11,16 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, GenerationConfig, PreTrainedModel
+from transformers import AutoConfig, AutoModelForSeq2SeqLM, GenerationConfig, PreTrainedModel
 
 from autogluon.timeseries.utils.warning_filters import set_loggers_level
 
-logger = logging.getLogger(__name__)
+from .base import BaseChronosPipeline, ForecastType
 
+logger = logging.getLogger("autogluon.timeseries.models.chronos")
 
-__all__ = ["ChronosConfig", "ChronosPipeline", "OptimizedChronosPipeline"]
+
+__all__ = ["ChronosConfig", "ChronosPipeline"]
 
 
 @dataclass
@@ -35,7 +37,7 @@ class ChronosConfig:
     pad_token_id: int
     eos_token_id: int
     use_eos_token: bool
-    model_type: Literal["causal", "seq2seq"]
+    model_type: Literal["seq2seq"]
     context_length: int
     prediction_length: int
     num_samples: int
@@ -279,18 +281,7 @@ class ChronosPretrainedModel(nn.Module):
         return preds.reshape(input_ids.size(0), num_samples, -1)
 
 
-def left_pad_and_stack_1D(tensors: List[torch.Tensor]):
-    max_len = max(len(c) for c in tensors)
-    padded = []
-    for c in tensors:
-        assert isinstance(c, torch.Tensor)
-        assert c.ndim == 1
-        padding = torch.full(size=(max_len - len(c),), fill_value=torch.nan, device=c.device)
-        padded.append(torch.concat((padding, c), dim=-1))
-    return torch.stack(padded)
-
-
-class ChronosPipeline:
+class ChronosPipeline(BaseChronosPipeline):
     """
     A ``ChronosPipeline`` uses the given tokenizer and model to forecast
     input time series.
@@ -308,21 +299,12 @@ class ChronosPipeline:
 
     tokenizer: ChronosTokenizer
     model: ChronosPretrainedModel
+    forecast_type: ForecastType = ForecastType.SAMPLES
 
     def __init__(self, tokenizer, model):
         self.tokenizer = tokenizer
         self.model = 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
-
     @torch.no_grad()
     def embed(self, context: Union[torch.Tensor, List[torch.Tensor]]) -> Tuple[torch.Tensor, Any]:
         """
@@ -363,7 +345,7 @@ class ChronosPipeline:
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_p: Optional[float] = None,
-        limit_prediction_length: bool = True,
+        limit_prediction_length: bool = False,
     ) -> torch.Tensor:
         """
         Get forecasts for the given time series.
@@ -442,42 +424,33 @@ class ChronosPipeline:
 
         return torch.cat(predictions, dim=-1)
 
-    @classmethod
-    def from_pretrained(cls, *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``.
-        """
-
-        config = AutoConfig.from_pretrained(*args, **kwargs)
-
-        assert hasattr(config, "chronos_config"), "Not a Chronos config file"
-
-        chronos_config = ChronosConfig(**config.chronos_config)
-
-        if chronos_config.model_type == "seq2seq":
-            inner_model = AutoModelForSeq2SeqLM.from_pretrained(*args, **kwargs)
-        else:
-            assert config.model_type == "causal"
-            inner_model = AutoModelForCausalLM.from_pretrained(*args, **kwargs)
-
-        return cls(
-            tokenizer=chronos_config.create_tokenizer(),
-            model=ChronosPretrainedModel(config=chronos_config, model=inner_model),
+    def predict_quantiles(
+        self,
+        context: torch.Tensor,
+        prediction_length: int,
+        quantile_levels: List[float],
+        num_samples: Optional[int] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        num_samples = num_samples or self.model.config.num_samples
+        prediction_samples = (
+            self.predict(
+                context,
+                prediction_length=prediction_length,
+                num_samples=num_samples,
+            )
+            .detach()
+            .cpu()
+            .swapaxes(1, 2)
         )
+        mean = prediction_samples.mean(axis=-1, keepdims=True)
+        quantiles = torch.quantile(
+            prediction_samples,
+            q=torch.tensor(quantile_levels, dtype=prediction_samples.dtype),
+            dim=-1,
+        ).permute(1, 2, 0)
 
-
-class OptimizedChronosPipeline(ChronosPipeline):
-    """A wrapper around the ChronosPipeline object for CPU-optimized model classes from
-    HuggingFace optimum.
-    """
-
-    dtypes = {
-        "bfloat16": torch.bfloat16,
-        "float32": torch.float32,
-        "float64": torch.float64,
-    }
+        return quantiles, mean
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
@@ -498,49 +471,40 @@ class OptimizedChronosPipeline(ChronosPipeline):
             config.chronos_config["context_length"] = context_length
         chronos_config = ChronosConfig(**config.chronos_config)
 
-        torch_dtype = kwargs.get("torch_dtype", "auto")
-        if torch_dtype != "auto" and isinstance(torch_dtype, str):
-            kwargs["torch_dtype"] = cls.dtypes[torch_dtype]
-
-        if chronos_config.model_type == "seq2seq":
-            if optimization_strategy is None:
-                inner_model = AutoModelForSeq2SeqLM.from_pretrained(*args, **kwargs)
-            else:
-                assert optimization_strategy in [
-                    "onnx",
-                    "openvino",
-                ], "optimization_strategy not recognized. Please provide one of `onnx` or `openvino`"
-                torch_dtype = kwargs.pop("torch_dtype", "auto")
-                if torch_dtype != "auto":
-                    logger.warning(
-                        f"\t`torch_dtype` will be ignored for optimization_strategy {optimization_strategy}"
+        assert chronos_config.model_type == "seq2seq"
+        if optimization_strategy is None:
+            inner_model = AutoModelForSeq2SeqLM.from_pretrained(*args, **kwargs)
+        else:
+            assert optimization_strategy in [
+                "onnx",
+                "openvino",
+            ], "optimization_strategy not recognized. Please provide one of `onnx` or `openvino`"
+            torch_dtype = kwargs.pop("torch_dtype", "auto")
+            if torch_dtype != "auto":
+                logger.warning(f"\t`torch_dtype` will be ignored for optimization_strategy {optimization_strategy}")
+
+            if optimization_strategy == "onnx":
+                try:
+                    from optimum.onnxruntime import ORTModelForSeq2SeqLM
+                except ImportError:
+                    raise ImportError(
+                        "Huggingface Optimum library must be installed with ONNX for using the `onnx` strategy"
                     )
 
-                if optimization_strategy == "onnx":
-                    try:
-                        from optimum.onnxruntime import ORTModelForSeq2SeqLM
-                    except ImportError:
-                        raise ImportError(
-                            "Huggingface Optimum library must be installed with ONNX for using the `onnx` strategy"
-                        )
-
-                    assert kwargs.pop("device_map", "cpu") in ["cpu", "auto"], "ONNX mode only available on the CPU"
-                    with set_loggers_level(regex=r"^optimum.*", level=logging.ERROR):
-                        inner_model = ORTModelForSeq2SeqLM.from_pretrained(*args, **{**kwargs, "export": True})
-                elif optimization_strategy == "openvino":
-                    try:
-                        from optimum.intel import OVModelForSeq2SeqLM
-                    except ImportError:
-                        raise ImportError(
-                            "Huggingface Optimum library must be installed with OpenVINO for using the `openvino` strategy"
-                        )
-                    with set_loggers_level(regex=r"^optimum.*", level=logging.ERROR):
-                        inner_model = OVModelForSeq2SeqLM.from_pretrained(
-                            *args, **{**kwargs, "device_map": "cpu", "export": True}
-                        )
-        else:
-            assert config.model_type == "causal"
-            inner_model = AutoModelForCausalLM.from_pretrained(*args, **kwargs)
+                assert kwargs.pop("device_map", "cpu") in ["cpu", "auto"], "ONNX mode only available on the CPU"
+                with set_loggers_level(regex=r"^optimum.*", level=logging.ERROR):
+                    inner_model = ORTModelForSeq2SeqLM.from_pretrained(*args, **{**kwargs, "export": True})
+            elif optimization_strategy == "openvino":
+                try:
+                    from optimum.intel import OVModelForSeq2SeqLM
+                except ImportError:
+                    raise ImportError(
+                        "Huggingface Optimum library must be installed with OpenVINO for using the `openvino` strategy"
+                    )
+                with set_loggers_level(regex=r"^optimum.*", level=logging.ERROR):
+                    inner_model = OVModelForSeq2SeqLM.from_pretrained(
+                        *args, **{**kwargs, "device_map": "cpu", "export": True}
+                    )
 
         return cls(
             tokenizer=chronos_config.create_tokenizer(),