use repo_id to download checkpoints

src/timesfm.py

@@ -24,6 +24,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
+from huggingface_hub import hf_hub_download
 from paxml import checkpoints
 from paxml import tasks_lib
 from praxis import base_hyperparams
@@ -222,7 +223,7 @@ class TimesFm:
 
   def load_from_checkpoint(
       self,
-      checkpoint_path: str,
+      repo_id: str = "google/timesfm-1.0-200m",
       checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX,
       step: int | None = None,
   ) -> None:
@@ -233,6 +234,9 @@ class TimesFm:
       checkpoint_type: type of PAX checkpoint
       step: step of the checkpoint to load. If `None`, load lastest checkpoint.
     """
+    # Download the checkpoint from Hugging Face Hub
+    checkpoint_path = hf_hub_download(repo_id)
+
     #  Initialize the model weights.
     self._logging("Constructing model weights.")
     start_time = time.time()