From e28239cb2ff3d7f2453d258350bf02816cbeaf12 Mon Sep 17 00:00:00 2001 From: misha-chertushkin Date: Wed, 25 Jun 2025 16:33:45 +0000 Subject: [PATCH 1/2] Added flag for safetensors loading --- pyproject.toml | 2 ++ src/finetuning/finetuning_example.py | 37 +++++++++++++++++++--------- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 09a0a18..0a6aac2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,8 @@ scikit-learn = ">=1.2.2" typer = ">=0.12.3" wandb = ">=0.17.5" absl-py = ">=1.4.0" +yfinance = "0.2.63" +safetensors = "^0.5.3" [tool.poetry.extras] pax = ["paxml", "lingvo", "jax", "jaxlib"] diff --git a/src/finetuning/finetuning_example.py b/src/finetuning/finetuning_example.py index 2c396d6..f1d76af 100644 --- a/src/finetuning/finetuning_example.py +++ b/src/finetuning/finetuning_example.py @@ -9,6 +9,7 @@ python script.py --training_mode=multi --gpu_ids=0,1,2 """ import os +from dataclasses import asdict from os import path from typing import Optional, Tuple @@ -19,11 +20,13 @@ import torch.multiprocessing as mp import yfinance as yf from absl import app, flags from huggingface_hub import snapshot_download +from safetensors.torch import load_file from torch.utils.data import Dataset from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams -from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder +from timesfm.pytorch_patched_decoder import (PatchedTimeSeriesDecoder, + TimesFMConfig) FLAGS = flags.FLAGS @@ -39,6 +42,11 @@ flags.DEFINE_list( "Comma-separated list of GPU IDs to use for multi-GPU training. Example: 0,1,2" ) +flags.DEFINE_string( + "local_model_path", + None, + "Path to a local .safetensors model file. If provided, overrides Hugging Face download." +) class TimeSeriesDataset(Dataset): """Dataset for time series data compatible with TimesFM.""" @@ -132,25 +140,32 @@ def prepare_datasets(series: np.ndarray, def get_model(load_weights: bool = False): device = "cuda" if torch.cuda.is_available() else "cpu" - repo_id = "google/timesfm-2.0-500m-pytorch" hparams = TimesFmHparams( backend=device, per_core_batch_size=32, horizon_len=128, num_layers=50, use_positional_embedding=False, - context_len= - 192, # Context length can be anything up to 2048 in multiples of 32 + context_len=192, ) - tfm = TimesFm(hparams=hparams, - checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id)) - - model = PatchedTimeSeriesDecoder(tfm._model_config) + if load_weights: - checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt") - loaded_checkpoint = torch.load(checkpoint_path, weights_only=True) + if FLAGS.local_model_path: + tfm_config = TimesFMConfig() + model = PatchedTimeSeriesDecoder(tfm_config) + loaded_checkpoint = load_file(FLAGS.local_model_path) + else: + repo_id = "google/timesfm-2.0-500m-pytorch" + tfm = TimesFm(hparams=hparams, + checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id)) + + tfm_config = tfm._model_config + model = PatchedTimeSeriesDecoder(tfm_config) + checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt") + loaded_checkpoint = torch.load(checkpoint_path, weights_only=True) + model.load_state_dict(loaded_checkpoint) - return model, hparams, tfm._model_config + return model, hparams, tfm_config def plot_predictions( From 6cf4f23675402f00f625ebfbd1942872ce4099bf Mon Sep 17 00:00:00 2001 From: misha-chertushkin Date: Wed, 25 Jun 2025 16:37:08 +0000 Subject: [PATCH 2/2] Remove yfinance --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0a6aac2..0bebccf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,6 @@ scikit-learn = ">=1.2.2" typer = ">=0.12.3" wandb = ">=0.17.5" absl-py = ">=1.4.0" -yfinance = "0.2.63" safetensors = "^0.5.3" [tool.poetry.extras]