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

Merge pull request #177 from andreramosfdc/allow_local_dir_download

Browse Source 
Enable local model snapshot downloads
This commit is contained in:
Yichen Zhou
2024-11-05 16:28:19 -08:00
committed by GitHub
parent e3d45d6568 c1e2094a10
commit 9b302ae794
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
+3 -2
View File
@@ -55,8 +55,9 @@ 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(
"torch_model.ckpt") snapshot_download(repo_id, local_dir=checkpoint.local_dir),
"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)
logging.info("Loading checkpoint from %s", checkpoint_path) logging.info("Loading checkpoint from %s", checkpoint_path)