wehub-resource-sync
97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
# 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
|
|
import types
|
|
|
|
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_torch_compile_wraps_forward(self):
|
|
"""Verifies that torch_compile=True compiles model.forward, not a no-op."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tfm = TimesFM_2p5_200M_torch(torch_compile=False)
|
|
tfm._save_pretrained(tmpdir)
|
|
|
|
tfm_compiled = TimesFM_2p5_200M_torch(torch_compile=True)
|
|
tfm_compiled.load_checkpoint(tmpdir)
|
|
|
|
# forward should be a compiled callable, not a plain bound method
|
|
assert not isinstance(tfm_compiled.model.forward, types.MethodType), (
|
|
"model.forward should be compiled after load_checkpoint with torch_compile=True"
|
|
)
|
|
|
|
def test_torch_no_compile_leaves_forward_unchanged(self):
|
|
"""Verifies that torch_compile=False leaves model.forward as a plain method."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tfm = TimesFM_2p5_200M_torch(torch_compile=False)
|
|
tfm._save_pretrained(tmpdir)
|
|
|
|
tfm_no_compile = TimesFM_2p5_200M_torch(torch_compile=False)
|
|
tfm_no_compile.load_checkpoint(tmpdir)
|
|
|
|
assert isinstance(tfm_no_compile.model.forward, types.MethodType), (
|
|
"model.forward should remain a plain bound method when torch_compile=False"
|
|
)
|
|
|
|
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
|