Merge pull request #312 from google-research/rajat-dev

Download using ModelHubMixin

README.md

@@ -59,8 +59,8 @@ pip install -e .
 ```python
 import numpy as np
 import timesfm
-model = timesfm.TimesFM_2p5_200M_torch()
+model = TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch")
-model.load_checkpoint()
+
 model.compile(
     timesfm.ForecastConfig(
         max_context=1024,

pyproject.toml

@@ -20,6 +20,10 @@ huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
 safetensors = ">=0.5.3"
 torch = { version = ">=2.0.0", extras = ["cuda"] }
 
+[tool.ruff]
+line-length = 88
+indent-width = 2
+
 
 [build-system]
 requires = ["poetry-core"]

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -11,24 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """TimesFM models."""
 
 import logging
 import math
 import os
-from typing import Sequence
+from pathlib import Path
+from typing import Dict, Optional, Sequence, Union
 
-import huggingface_hub
 import numpy as np
-from safetensors.torch import load_file
 import torch
+from huggingface_hub import ModelHubMixin, hf_hub_download
+from safetensors.torch import load_file
 from torch import nn
 
 from .. import configs
-from ..torch import dense
+from ..torch import dense, transformer, util
-from ..torch import transformer
-from ..torch import util
 from . import timesfm_2p5_base
 
 revin = util.revin
@@ -56,10 +54,12 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
 
     # Layers.
     self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
-    self.stacked_xf = nn.ModuleList([
+    self.stacked_xf = nn.ModuleList(
+      [
         transformer.Transformer(self.config.stacked_transformers.transformer)
         for _ in range(self.x)
-    ])
+      ]
+    )
     self.output_projection_point = dense.ResidualBlock(
       self.config.output_projection_point
     )
@@ -144,12 +144,8 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
 
       decode_caches = [
         util.DecodeCache(
-              next_index=torch.zeros(
+          next_index=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device),
-                  batch_size, dtype=torch.int32, device=inputs.device
+          num_masked=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device),
-              ),
-              num_masked=torch.zeros(
-                  batch_size, dtype=torch.int32, device=inputs.device
-              ),
           key=torch.zeros(
             batch_size,
             decode_cache_size,
@@ -168,9 +164,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
         for _ in range(self.x)
       ]
 
-      normed_inputs = revin(
+      normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False)
-          patched_inputs, context_mu, context_sigma, reverse=False
-      )
       normed_inputs = torch.where(patched_masks, 0.0, normed_inputs)
       (_, _, normed_outputs, normed_quantile_spread), decode_caches = self(
         normed_inputs, patched_masks, decode_caches
@@ -180,9 +174,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
         (batch_size, -1, self.o, self.q),
       )
       renormed_quantile_spread = torch.reshape(
-          revin(
+        revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
-              normed_quantile_spread, context_mu, context_sigma, reverse=True
-          ),
         (batch_size, -1, self.os, self.q),
       )[:, -1, ...]
 
@@ -208,9 +200,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
         new_mu = torch.stack(new_mus, dim=1)
         new_sigma = torch.stack(new_sigmas, dim=1)
 
-        new_normed_input = revin(
+        new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False)
-            new_patched_input, new_mu, new_sigma, reverse=False
-        )
         (_, _, new_normed_output, _), decode_caches = self(
           new_normed_input, new_mask, decode_caches
         )
@@ -254,9 +244,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
         input_t = torch.cat(
           [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0
         )
-        mask = torch.cat(
+        mask = torch.cat([torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0)
-            [torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0
-        )
       input_t = input_t[None, ...]
       mask = mask[None, ...]
       t_pf, _, t_ar = self.decode(horizon, input_t, mask)
@@ -269,38 +257,58 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
     return outputs
 
 
-class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
+class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
   """PyTorch implementation of TimesFM 2.5 with 200M parameters."""
 
   model: nn.Module = TimesFM_2p5_200M_torch_module()
 
-  def load_checkpoint(
+  @classmethod
-      self,
+  def _from_pretrained(
+    cls,
     *,
-      path: str | None = None,
+    model_id: str,
-      hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch",
+    revision: Optional[str],
-  ) -> None:
+    cache_dir: Optional[Union[str, Path]],
-    """Loads a PyTorch safetensors TimesFM model.
+    force_download: bool,
-
+    proxies: Optional[Dict],
-    Args:
+    resume_download: Optional[bool],
-        path: Path to a local checkpoint. If not provided, will try to download
+    local_files_only: bool,
-          from the default Hugging Face repo.
+    token: Optional[str],
-        hf_repo_id: If provided, will download from the specified Hugging Face
+    **model_kwargs,
-          repo instead.
+  ):
     """
-    if path:
+    Loads a PyTorch safetensors TimesFM model from a local path or the Hugging
-      pass
+    Face Hub. This method is the backend for the `from_pretrained` class
-    elif hf_repo_id:
+    method provided by `ModelHubMixin`.
-      logging.info(
+    """
-          "Downloading checkpoint from Hugging Face repo %s", hf_repo_id
+    # Create an instance of the model wrapper class.
-      )
+    instance = cls(**model_kwargs)
-      path = os.path.join(
+
-          huggingface_hub.snapshot_download(hf_repo_id), "model.safetensors"
+    # Determine the path to the model weights.
-      )
+    model_file_path = ""
-      logging.info("Loading checkpoint from: %s", path)
+    if os.path.isdir(model_id):
+      logging.info("Loading checkpoint from local directory: %s", model_id)
+      model_file_path = os.path.join(model_id, "model.safetensors")
+      if not os.path.exists(model_file_path):
+        raise FileNotFoundError(f"model.safetensors not found in directory {model_id}")
     else:
-      raise ValueError("Either path or hf_repo_id must be provided.")
+      logging.info("Downloading checkpoint from Hugging Face repo %s", model_id)
-    self.model.load_checkpoint(path)
+      model_file_path = hf_hub_download(
+        repo_id=model_id,
+        filename="model.safetensors",
+        revision=revision,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        proxies=proxies,
+        resume_download=resume_download,
+        token=token,
+        local_files_only=local_files_only,
+      )
+
+    logging.info("Loading checkpoint from: %s", model_file_path)
+    # Load the weights into the model.
+    instance.model.load_checkpoint(model_file_path)
+    return instance
 
   def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None:
     """Attempts to compile the model for fast decoding.
@@ -326,8 +334,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
         "When compiling, max context needs to be multiple of the patch size"
         " %d. Using max context = %d instead.",
         self.model.p,
-          new_context := math.ceil(fc.max_context / self.model.p)
+        new_context := math.ceil(fc.max_context / self.model.p) * self.model.p,
-          * self.model.p,
       )
       fc.max_context = new_context
     if fc.max_horizon % self.model.o != 0:
@@ -335,8 +342,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
         "When compiling, max horizon needs to be multiple of the output patch"
         " size %d. Using max horizon = %d instead.",
         self.model.o,
-          new_horizon := math.ceil(fc.max_horizon / self.model.o)
+        new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o,
-          * self.model.o,
       )
       fc.max_horizon = new_horizon
     if fc.max_context + fc.max_horizon > self.model.config.context_limit:
@@ -347,16 +353,14 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
       )
     if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os):
       raise ValueError(
-          "Continuous quantile head is not supported for horizons >"
+        f"Continuous quantile head is not supported for horizons > {self.model.os}."
-          f" {self.model.os}."
       )
     self.forecast_config = fc
 
     def _compiled_decode(horizon, inputs, masks):
       if horizon > fc.max_horizon:
         raise ValueError(
-            "Horizon must be less than the max horizon."
+          f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}."
-            f" {horizon} > {fc.max_horizon}."
         )
 
       inputs = torch.Tensor(np.array(inputs)).to(self.model.device)
@@ -383,9 +387,8 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
         to_cat.append(ar_outputs.reshape(batch_size, -1, self.model.q))
       full_forecast = torch.cat(to_cat, dim=1)
 
-      flip_quantile_fn = lambda x: torch.cat(
+      def flip_quantile_fn(x):
-          [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1
+        return torch.cat([x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1)
-      )
 
       if fc.force_flip_invariance:
         flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
@@ -395,9 +398,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
         flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs)
         to_cat = [flipped_pf_outputs[:, -1, ...]]
         if flipped_ar_outputs is not None:
-          to_cat.append(
+          to_cat.append(flipped_ar_outputs.reshape(batch_size, -1, self.model.q))
-              flipped_ar_outputs.reshape(batch_size, -1, self.model.q)
-          )
         flipped_full_forecast = torch.cat(to_cat, dim=1)
         quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2
         pf_outputs = (pf_outputs - flipped_pf_outputs) / 2