freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

use repo_id to download checkpoints

Browse Source
This commit is contained in:
Kashif Rasul
2024-05-10 10:48:49 +02:00
parent e6fb5140f4
commit afeb84f77a
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/timesfm.py
+5 -1
View File
@@ -24,6 +24,7 @@ import jax
import jax.numpy as jnp import jax.numpy as jnp
import numpy as np import numpy as np
import pandas as pd import pandas as pd
from huggingface_hub import hf_hub_download
from paxml import checkpoints from paxml import checkpoints
from paxml import tasks_lib from paxml import tasks_lib
from praxis import base_hyperparams from praxis import base_hyperparams
@@ -222,7 +223,7 @@ class TimesFm:
def load_from_checkpoint( def load_from_checkpoint(
self, self,
checkpoint_path: str, repo_id: str = "google/timesfm-1.0-200m",
checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX, checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX,
step: int | None = None, step: int | None = None,
) -> None: ) -> None:
@@ -233,6 +234,9 @@ class TimesFm:
checkpoint_type: type of PAX checkpoint checkpoint_type: type of PAX checkpoint
step: step of the checkpoint to load. If `None`, load lastest 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. # Initialize the model weights.
self._logging("Constructing model weights.") self._logging("Constructing model weights.")
start_time = time.time() start_time = time.time()