From afeb84f77ac86686389337e528caecfbc496ad75 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Fri, 10 May 2024 10:48:49 +0200
Subject: [PATCH] use repo_id to download checkpoints

---
 src/timesfm.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/timesfm.py b/src/timesfm.py
index b795450..56456f5 100644
--- a/src/timesfm.py
+++ b/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()