From 2f1625c208cf99bd4c9f6650d54796d8f88486f8 Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Sun, 7 Jun 2026 17:28:32 +0000 Subject: [PATCH] Fix model loading issues and forecast_naive slicing bug in TimesFM 2.5 - Allow model wrapper constructors (__init__) to accept and ignore extra keyword arguments (e.g. proxies) passed by huggingface_hub during from_pretrained. - Implement load_checkpoint for TimesFM_2p5_200M_torch and TimesFM_2p5_200M_flax to restore weights from local paths. - Fix slicing bug in PyTorch's forecast_naive to correctly slice the time/horizon dimension ([:, :horizon, :]) instead of quantiles. - Add unit tests in tests/test_model_loading.py covering local checkpoint loading, hub compatibility, and prediction shape correctness. --- src/timesfm/timesfm_2p5/timesfm_2p5_flax.py | 20 ++++-- src/timesfm/timesfm_2p5/timesfm_2p5_torch.py | 18 +++++- tests/test_model_loading.py | 68 ++++++++++++++++++++ 3 files changed, 100 insertions(+), 6 deletions(-) create mode 100644 tests/test_model_loading.py diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py index 0f3b150..b823af3 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py @@ -447,6 +447,21 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5): model: nnx.Module = TimesFM_2p5_200M_flax_module() + def __init__(self, **kwargs): + self.model = TimesFM_2p5_200M_flax_module() + + def load_checkpoint(self, path: str): + """Loads a TimesFM model from a checkpoint.""" + if os.path.isdir(path): + model_file_path = path + else: + model_file_path = os.path.dirname(path) + + checkpointer = ocp.StandardCheckpointer() + graph, state = nnx.split(self.model) + state = checkpointer.restore(model_file_path, state) + self.model = nnx.merge(graph, state) + @classmethod def from_pretrained( cls, @@ -485,10 +500,7 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5): ) logging.info("Loading checkpoint from: %s", model_file_path) - checkpointer = ocp.StandardCheckpointer() - graph, state = nnx.split(instance.model) - state = checkpointer.restore(model_file_path, state) - instance.model = nnx.merge(graph, state) + instance.load_checkpoint(model_file_path) return instance def compile( diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 3e7c9f1..e41ad5a 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -257,7 +257,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module): to_concat = [t_pf[:, -1, ...]] if t_ar is not None: to_concat.append(t_ar.reshape(1, -1, self.q)) - torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon] + torch_forecast = torch.cat(to_concat, dim=1)[:, :horizon, :] torch_forecast = torch_forecast.squeeze(0) outputs.append(torch_forecast.detach().cpu().numpy()) return outputs @@ -283,12 +283,26 @@ class TimesFM_2p5_200M_torch( self, torch_compile: bool = True, config: Optional[dict] = None, + **kwargs, ): self.model = TimesFM_2p5_200M_torch_module() self.torch_compile = torch_compile if config is not None: self._hub_mixin_config = config + def load_checkpoint(self, path: str, **kwargs): + """Loads a TimesFM model from a checkpoint directory or file.""" + if os.path.isdir(path): + model_file_path = os.path.join(path, self.WEIGHTS_FILENAME) + if not os.path.exists(model_file_path): + raise FileNotFoundError( + f"{self.WEIGHTS_FILENAME} not found in directory {path}" + ) + else: + model_file_path = path + + self.model.load_checkpoint(model_file_path, **kwargs) + @classmethod def _from_pretrained( cls, @@ -333,7 +347,7 @@ class TimesFM_2p5_200M_torch( logging.info("Loading checkpoint from: %s", model_file_path) # Load the weights into the model. - instance.model.load_checkpoint( + instance.load_checkpoint( model_file_path, torch_compile=instance.torch_compile ) return instance diff --git a/tests/test_model_loading.py b/tests/test_model_loading.py new file mode 100644 index 0000000..998cab4 --- /dev/null +++ b/tests/test_model_loading.py @@ -0,0 +1,68 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for loading TimesFM 2.5 models.""" + +import os +import tempfile + +from timesfm.timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch +from timesfm.timesfm_2p5.timesfm_2p5_flax import TimesFM_2p5_200M_flax + + +class TestModelLoading: + """Tests to verify model instantiation, loading, and compatibility.""" + + def test_torch_load_checkpoint_and_from_pretrained_local(self): + """Verifies that PyTorch load_checkpoint and from_pretrained work locally.""" + # 1. Instantiate the model wrapper with compilation disabled + tfm = TimesFM_2p5_200M_torch(torch_compile=False) + + with tempfile.TemporaryDirectory() as tmpdir: + # 2. Save the model's randomly-initialized weights + tfm._save_pretrained(tmpdir) + + # Verify weights file is written + weights_path = os.path.join(tmpdir, "model.safetensors") + assert os.path.exists(weights_path) + + # 3. Verify that load_checkpoint works from the temp directory path + tfm2 = TimesFM_2p5_200M_torch(torch_compile=False) + tfm2.load_checkpoint(tmpdir, torch_compile=False) + + # 4. Verify that from_pretrained works with a local directory path + # and accepts/ignores extra kwargs (like proxies) without raising TypeError + tfm3 = TimesFM_2p5_200M_torch.from_pretrained( + tmpdir, + torch_compile=False, + proxies={"http": "http://dummy.proxy"}, + custom_kwarg="dummy_value", + ) + assert tfm3 is not None + assert not tfm3.torch_compile + + # 5. Run a simple prediction step to verify the loaded model performs forward pass + import numpy as np + inputs = [np.random.randn(32)] + forecasts = tfm3.model.forecast_naive(horizon=10, inputs=inputs) + assert len(forecasts) == 1 + assert forecasts[0].shape == (10, 10) + + def test_flax_model_init_kwargs(self): + """Verifies that Flax model wrapper constructor accepts arbitrary kwargs.""" + tfm = TimesFM_2p5_200M_flax( + proxies={"http": "http://dummy.proxy"}, + custom_kwarg="dummy_value", + ) + assert tfm is not None