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

Enable model snapshot downloads to be saved in a specified local directory.

Browse Source
This commit is contained in:
andreramosfdc
2024-10-30 18:08:17 -03:00
committed by andre_ramos
parent e3d45d6568
commit c1e2094a10
2 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/timesfm/timesfm_base.py
+1
View File
@@ -130,6 +130,7 @@ class TimesFmCheckpoint:
huggingface_repo_id: str | None = None huggingface_repo_id: str | None = None
type: Any = None type: Any = None
step: int | None = None step: int | None = None
local_dir: str | None = None
class TimesFmBase: class TimesFmBase:
src/timesfm/timesfm_torch.py
+2 -1
View File
@@ -55,7 +55,8 @@ class TimesFmTorch(timesfm_base.TimesFmBase):
checkpoint_path = checkpoint.path checkpoint_path = checkpoint.path
repo_id = checkpoint.huggingface_repo_id repo_id = checkpoint.huggingface_repo_id
if checkpoint_path is None: if checkpoint_path is None:
checkpoint_path = path.join(snapshot_download(repo_id), checkpoint_path = path.join(
snapshot_download(repo_id, local_dir=checkpoint.local_dir),
"torch_model.ckpt") "torch_model.ckpt")
self._model = ppd.PatchedTimeSeriesDecoder(self._model_config) self._model = ppd.PatchedTimeSeriesDecoder(self._model_config)
loaded_checkpoint = torch.load(checkpoint_path, weights_only=True) loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)