From afeb84f77ac86686389337e528caecfbc496ad75 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 10 May 2024 10:48:49 +0200 Subject: [PATCH 1/5] 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() From d5f55b08256633eb2ed7a6f2708568ccc9e68abb Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 10 May 2024 11:32:53 +0200 Subject: [PATCH 2/5] use snapshot_download --- src/timesfm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/timesfm.py b/src/timesfm.py index 56456f5..23635eb 100644 --- a/src/timesfm.py +++ b/src/timesfm.py @@ -24,7 +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 huggingface_hub import snapshot_download from paxml import checkpoints from paxml import tasks_lib from praxis import base_hyperparams @@ -235,7 +235,7 @@ class TimesFm: 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) + checkpoint_path = snapshot_download(repo_id) # Initialize the model weights. self._logging("Constructing model weights.") From f04d4d637d5e5db8b5ef8b602ff5038ab8f1729e Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 10 May 2024 11:39:38 +0200 Subject: [PATCH 3/5] update readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 916a37d..73e5884 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,7 @@ tfm = timesfm.TimesFm( model_dims=1280, backend=, ) -tfm.load_from_checkpoint() +tfm.load_from_checkpoint(repo_id="google/timesfm-1.0-200m") ``` Note that the four parameters are fixed to load the 200m model From 4a2b0ae837e01e5e50006b3a2b3d5e27008dde7f Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 10 May 2024 12:01:42 +0200 Subject: [PATCH 4/5] fix checkpoint_path --- src/timesfm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/timesfm.py b/src/timesfm.py index 23635eb..8ae975c 100644 --- a/src/timesfm.py +++ b/src/timesfm.py @@ -16,6 +16,7 @@ import logging import multiprocessing +from os import path import time from typing import Any, Literal, Sequence @@ -235,7 +236,7 @@ class TimesFm: step: step of the checkpoint to load. If `None`, load lastest checkpoint. """ # Download the checkpoint from Hugging Face Hub - checkpoint_path = snapshot_download(repo_id) + checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints") # Initialize the model weights. self._logging("Constructing model weights.") From 2a83a0f82d898e995c886ff06bcb0488e93ff51d Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 10 May 2024 12:23:25 +0200 Subject: [PATCH 5/5] add back checkpoint_path --- src/timesfm.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/timesfm.py b/src/timesfm.py index 8ae975c..2887aad 100644 --- a/src/timesfm.py +++ b/src/timesfm.py @@ -18,7 +18,7 @@ import logging import multiprocessing from os import path import time -from typing import Any, Literal, Sequence +from typing import Any, Literal, Optional, Sequence import einshape as es import jax @@ -224,6 +224,7 @@ class TimesFm: def load_from_checkpoint( self, + checkpoint_path: Optional[str] = None, repo_id: str = "google/timesfm-1.0-200m", checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX, step: int | None = None, @@ -231,12 +232,14 @@ class TimesFm: """Loads a checkpoint and compiles the decoder. Args: - checkpoint_path: path to the checkpoint directory. + checkpoint_path: Optional path to the checkpoint directory. + repo_id: Hugging Face Hub repo id. 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 = path.join(snapshot_download(repo_id), "checkpoints") + # Download the checkpoint from Hugging Face Hub if not given + if checkpoint_path is None: + checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints") # Initialize the model weights. self._logging("Constructing model weights.")