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

Add troubleshooting section to readme file

Browse Source
This commit is contained in:
Funto-Adeyemi
2023-11-11 09:39:14 +00:00
parent 95e81c3539
commit 04a5a99b43
15 changed files with 923 additions and 979 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.style.yapf
+4
View File
@@ -0,0 +1,4 @@
[style]
based_on_style = google
indent_width = 2
spaces_before_comment = 2
experiments/baselines/timegpt_pipeline.py
+11 -13
View File
@@ -34,9 +34,8 @@ def get_seasonality(freq: str) -> int:
return _get_seasonality(freq, seasonalities={"D": 7})
def maybe_convert_col_to_datetime(
df: pd.DataFrame, col_name: str
) -> pd.DataFrame:
def maybe_convert_col_to_datetime(df: pd.DataFrame,
col_name: str) -> pd.DataFrame:
if not pd.api.types.is_datetime64_any_dtype(df[col_name]):
df = df.copy()
df[col_name] = pd.to_datetime(df[col_name])
@@ -64,13 +63,14 @@ def zero_pad_time_series(df, freq, min_length=36):
end=start_date,
periods=min_length - len(subset) + 1,
freq=freq, # 'MS' for month start
)[
:-1
] # Exclude the start_date itself
)[:-1] # Exclude the start_date itself
# 2c. Create padding data
padding_df = pd.DataFrame(
{"ds": padding_dates, "unique_id": unique_id, "y": 0} # Zero padding
padding_df = pd.DataFrame({
"ds": padding_dates,
"unique_id": unique_id,
"y": 0
} # Zero padding
)
# 2d. Combine original and padding data, and append to the list
@@ -121,8 +121,7 @@ class Forecaster:
for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
if len(valid.columns) > 3:
raise NotImplementedError(
"Cross validation with exogenous variables is not yet supported."
)
"Cross validation with exogenous variables is not yet supported.")
y_pred = self.forecast(
df=train,
h=h,
@@ -138,8 +137,7 @@ class Forecaster:
raise ValueError(
"Cross validation result produced less results than expected."
" Please verify that the frequency parameter (freq) matches your"
" series' and that there aren't any missing periods."
)
" series' and that there aren't any missing periods.")
results.append(result)
out = vertical_concat(results)
out = drop_index_if_pandas(out)
@@ -203,7 +201,7 @@ class TimeGPT(Forecaster):
all_unique_ids = df["unique_id"].unique()
all_fcst_df = []
for i in range(0, len(all_unique_ids), chunk_size):
chunk_ids = all_unique_ids[i : i + chunk_size]
chunk_ids = all_unique_ids[i:i + chunk_size]
chunk_df = df[df["unique_id"].isin(chunk_ids)]
fct_chunk_df = client.forecast(
df=chunk_df,
experiments/extended_benchmarks/run_timegpt.py
+3 -6
View File
@@ -11,7 +11,6 @@
# 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.
"""Evaluation script for timegpt."""
import os
@@ -25,7 +24,6 @@ import pandas as pd
from ..baselines.timegpt_pipeline import run_timegpt
from .utils import ExperimentHandler
dataset_names = [
"m1_monthly",
"m1_quarterly",
@@ -63,7 +61,6 @@ _MODEL_NAME = flags.DEFINE_string(
)
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
QUANTILES = list(np.arange(1, 10) / 10.0)
@@ -90,9 +87,9 @@ def main():
)
time_df = pd.DataFrame({"time": [total_time], "model": model_name})
fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name)
results = exp.evaluate_from_predictions(
models=[model_name], fcsts_df=fcsts_df, times_df=time_df
)
results = exp.evaluate_from_predictions(models=[model_name],
fcsts_df=fcsts_df,
times_df=time_df)
print(results, flush=True)
results_list.append(results)
results_full = pd.concat(results_list)
experiments/extended_benchmarks/run_timesfm.py
-1
View File
@@ -54,7 +54,6 @@ dataset_names = [
"hospital",
]
context_dict_v2 = {}
context_dict_v1 = {
experiments/extended_benchmarks/utils.py
+23 -35
View File
@@ -11,7 +11,6 @@
# 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.
"""Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py."""
from functools import partial
@@ -46,11 +45,9 @@ def quantile_loss(
target_col: str = "y",
) -> pd.DataFrame:
delta_y = df[models].sub(df[target_col], axis=0)
res = (
np.maximum(q * delta_y, (q - 1) * delta_y)
.groupby(df[id_col], observed=True)
.mean()
)
res = (np.maximum(q * delta_y,
(q - 1) * delta_y).groupby(df[id_col],
observed=True).mean())
res.index.name = id_col
res = res.reset_index()
return res
@@ -66,10 +63,8 @@ class ExperimentHandler:
models_dir: str = "./models",
):
if dataset not in gluonts_datasets:
raise Exception(
f"dataset {dataset} not found in gluonts "
f"available datasets: {', '.join(gluonts_datasets)}"
)
raise Exception(f"dataset {dataset} not found in gluonts "
f"available datasets: {', '.join(gluonts_datasets)}")
self.dataset = dataset
self.quantiles = quantiles
self.level = self._transform_quantiles_to_levels(quantiles)
@@ -80,10 +75,8 @@ class ExperimentHandler:
gluonts_dataset = get_dataset(self.dataset)
self.horizon = gluonts_dataset.metadata.prediction_length
if self.horizon is None:
raise Exception(
f"horizon not found for dataset {self.dataset} "
"experiment cannot be run"
)
raise Exception(f"horizon not found for dataset {self.dataset} "
"experiment cannot be run")
self.freq = gluonts_dataset.metadata.freq
# get_seasonality() returns 1 for freq='D', override this to 7. This significantly improves the accuracy of
# statistical models on datasets like m5/nn5_daily. The models like AutoARIMA/AutoETS can still set
@@ -122,8 +115,7 @@ class ExperimentHandler:
@staticmethod
def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]:
level = [
int(100 - 200 * q) for q in quantiles if q < 0.5
level = [int(100 - 200 * q) for q in quantiles if q < 0.5
] # in this case mean=mediain
level = sorted(list(set(level)))
return level
@@ -153,9 +145,8 @@ class ExperimentHandler:
last_n: int | None = None,
) -> pd.DataFrame:
with multiprocessing.Pool(os.cpu_count()) as pool: # Create a process pool
results = pool.map(
parallel_transform, zip(gluonts_dataset, repeat(last_n))
)
results = pool.map(parallel_transform, zip(gluonts_dataset,
repeat(last_n)))
df = pd.concat(results)
df = df.reset_index(drop=True)
return df
@@ -177,9 +168,8 @@ class ExperimentHandler:
def save_dataframe(self, df: pd.DataFrame, file_name: str):
df.to_csv(f"{self.results_dir}/{file_name}", index=False)
def save_results(
self, fcst_df: pd.DataFrame, total_time: float, model_name: str
):
def save_results(self, fcst_df: pd.DataFrame, total_time: float,
model_name: str):
self.save_dataframe(
fcst_df,
f"{model_name}-{self.dataset}-fcst.csv",
@@ -215,23 +205,21 @@ class ExperimentHandler:
times_df = []
for model in models:
fcst_method_df = pd.read_csv(
f"{self.results_dir}/{model}-{self.dataset}-fcst.csv"
).set_index(["unique_id", "ds"])
f"{self.results_dir}/{model}-{self.dataset}-fcst.csv").set_index(
["unique_id", "ds"])
fcsts_df.append(fcst_method_df)
time_method_df = pd.read_csv(
f"{self.results_dir}/{model}-{self.dataset}-time.csv"
)
f"{self.results_dir}/{model}-{self.dataset}-time.csv")
times_df.append(time_method_df)
fcsts_df = pd.concat(fcsts_df, axis=1).reset_index()
fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"])
times_df = pd.concat(times_df)
return self.evaluate_from_predictions(
models=models, fcsts_df=fcsts_df, times_df=times_df
)
return self.evaluate_from_predictions(models=models,
fcsts_df=fcsts_df,
times_df=times_df)
def evaluate_from_predictions(
self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame
) -> pd.DataFrame:
def evaluate_from_predictions(self, models: List[str], fcsts_df: pd.DataFrame,
times_df: pd.DataFrame) -> pd.DataFrame:
test_df = self.test_df
train_df = self.train_df
test_df = test_df.merge(fcsts_df, how="left")
@@ -262,9 +250,9 @@ class ExperimentHandler:
eval_prob_df["metric"] = "scaled_crps"
eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True)
eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index()
eval_df = eval_df.melt(
id_vars="metric", value_name="value", var_name="model"
)
eval_df = eval_df.melt(id_vars="metric",
value_name="value",
var_name="model")
times_df.insert(0, "metric", "time")
times_df = times_df.rename(columns={"time": "value"})
eval_df = pd.concat([eval_df, times_df])
peft/finetune.py
+69 -72
View File
@@ -11,7 +11,6 @@
# 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.
"""
Finetune pipeline.
"""
@@ -39,13 +38,11 @@ from timesfm import TimesFm, data_loader, patched_decoder
NestedMap = py_utils.NestedMap
warnings.filterwarnings("ignore")
cmdstanpy_logger = logging.getLogger("cmdstanpy")
absl_logger = logging.getLogger("absl")
cmdstanpy_logger.disabled = True
absl_logger.disabled = True
"""
TimesFM model config. These are fixed since pre-training was done
with this configuration.
@@ -62,20 +59,24 @@ RANDOM_SEED = 1234
def finetune(
*,
model_name: Annotated[
str, typer.Option(help="Specify the name of the huggingface model.")
] = "google/timesfm-1.0-200m",
model_name: Annotated[str,
typer.Option(
help="Specify the name of the huggingface model."
)] = "google/timesfm-1.0-200m",
checkpoint_path: Annotated[
str, typer.Option(help="The path to the local model checkpoint.")
] = None,
datetime_col: Annotated[str, typer.Option(help="Column having datetime.")] = "ds",
ts_cols: Annotated[
list[str], typer.Option(help="Columns of time-series features.")
] = [],
normalize: Annotated[
bool, typer.Option(help="Normalize data for eval or not")
] = True,
context_len: Annotated[int, typer.Option(help="Length of the context window")],
str,
typer.Option(help="The path to the local model checkpoint.")] = None,
datetime_col: Annotated[str,
typer.Option(
help="Column having datetime.")] = "ds",
ts_cols: Annotated[list[str],
typer.Option(
help="Columns of time-series features.")] = [],
normalize: Annotated[bool,
typer.Option(
help="Normalize data for eval or not")] = True,
context_len: Annotated[int,
typer.Option(help="Length of the context window")],
horizon_len: Annotated[int, typer.Option(help="Prediction length.")],
freq: Annotated[
str,
@@ -87,67 +88,66 @@ def finetune(
data_path: Annotated[str, typer.Option(help="Path to dataset csv")],
boundaries: Annotated[
Tuple[int, int, int],
typer.Option(
help="boundaries of dataset to train, val, test",
),
typer.Option(help="boundaries of dataset to train, val, test",),
] = (0, 0, 0),
backend: Annotated[str, typer.Option(help="Backend device: cpu, gpu, tpu")],
backend: Annotated[str,
typer.Option(help="Backend device: cpu, gpu, tpu")],
batch_size: Annotated[
int, typer.Option(help="Batch size for the randomly sampled batch")
] = 16,
int,
typer.Option(help="Batch size for the randomly sampled batch")] = 16,
num_epochs: Annotated[int, typer.Option(help="Number of epochs")],
learning_rate: Annotated[float, typer.Option(help="adam optimizer learning rate")],
adam_epsilon: Annotated[float, typer.Option(help="adam optimizer epsilon")],
adam_clip_threshold: Annotated[
float, typer.Option(help="adam optimizer clip threshold")
],
cos_initial_decay_value: Annotated[
float, typer.Option(help="cosine initial decay value")
],
cos_final_decay_value: Annotated[
float, typer.Option(help="cosine final decay value")
],
cos_decay_steps: Annotated[int, typer.Option(help="Number of cosine decay steps")],
ema_decay: Annotated[float, typer.Option(help="Exponential moving average decay")],
learning_rate: Annotated[float,
typer.Option(help="adam optimizer learning rate")],
adam_epsilon: Annotated[float,
typer.Option(help="adam optimizer epsilon")],
adam_clip_threshold: Annotated[float,
typer.Option(
help="adam optimizer clip threshold")],
cos_initial_decay_value: Annotated[float,
typer.Option(
help="cosine initial decay value")],
cos_final_decay_value: Annotated[float,
typer.Option(
help="cosine final decay value")],
cos_decay_steps: Annotated[int,
typer.Option(
help="Number of cosine decay steps")],
ema_decay: Annotated[float,
typer.Option(help="Exponential moving average decay")],
early_stop_patience: Annotated[
int, typer.Option(..., help="Early stopping patience")
] = 5,
int, typer.Option(..., help="Early stopping patience")] = 5,
use_lora: Annotated[
bool,
typer.Option(
help="Train low rank adapters for stacked transformer block",
),
typer.
Option(help="Train low rank adapters for stacked transformer block",),
] = False,
lora_rank: Annotated[
int,
typer.Option(
help="LoRA Rank",
),
typer.Option(help="LoRA Rank",),
] = 8,
lora_target_modules: Annotated[
str,
typer.Option(
help="LoRA target modules of the transformer block. Allowed values: [all, attention, mlp]"
help=
"LoRA target modules of the transformer block. Allowed values: [all, attention, mlp]"
),
] = "all",
use_dora: Annotated[
bool,
typer.Option(
help="Apply DoRA strategy along with LoRA.",
),
typer.Option(help="Apply DoRA strategy along with LoRA.",),
] = False,
use_linear_probing: Annotated[
bool,
typer.Option(
help="Linear Probing. Train only input/output and embedding params. Freeze params in stack transformer block.",
help=
"Linear Probing. Train only input/output and embedding params. Freeze params in stack transformer block.",
),
] = False,
checkpoint_dir: Annotated[
str, typer.Option(help="Checkpoint directory")
] = "./checkpoints",
wandb_project: Annotated[
str, typer.Option(help="Weights & Biases project name")
] = "google_timesfm_finetune",
str, typer.Option(help="Checkpoint directory")] = "./checkpoints",
wandb_project: Annotated[str,
typer.Option(help="Weights & Biases project name"
)] = "google_timesfm_finetune",
) -> None:
key = jax.random.PRNGKey(seed=RANDOM_SEED)
wandb.init(project=wandb_project, config=locals())
@@ -261,9 +261,7 @@ def finetune(
task_p = tasks_lib.SingleTask(
name="ts-learn",
model=model,
train=tasks_lib.SingleTask.Train(
learner=build_learner(),
),
train=tasks_lib.SingleTask.Train(learner=build_learner(),),
)
task_p.model.ici_mesh_shape = [1, 1, 1]
@@ -296,18 +294,19 @@ def finetune(
checkpoint_type=checkpoint_types.CheckpointType.GDA,
)
jax_model_states.mdl_vars["params"]["core_layer"] = tfm._train_state.mdl_vars[
"params"
]
"params"]
gc.collect()
jax_task = task_p
def train_step(states, prng_key, inputs):
return trainer_lib.train_step_single_learner(jax_task, states, prng_key, inputs)
return trainer_lib.train_step_single_learner(jax_task, states, prng_key,
inputs)
def eval_step(states, prng_key, inputs):
states = states.to_eval_state()
return trainer_lib.eval_step_single_learner(jax_task, states, prng_key, inputs)
return trainer_lib.eval_step_single_learner(jax_task, states, prng_key,
inputs)
key, train_key, eval_key = jax.random.split(key, 3)
train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())
@@ -319,6 +318,7 @@ def finetune(
replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)
def reshape_batch_for_pmap(batch, num_devices):
def _reshape(input_tensor):
bsize = input_tensor.shape[0]
residual_shape = list(input_tensor.shape[1:])
@@ -341,8 +341,7 @@ def finetune(
tbatch = process_train_batch(batch)
tbatch = reshape_batch_for_pmap(tbatch, num_devices)
replicated_jax_states, step_fun_out = p_train_step(
replicated_jax_states, train_prng_seed, tbatch
)
replicated_jax_states, train_prng_seed, tbatch)
train_losses.append(step_fun_out.loss[0])
wandb.log({"train_step_loss": step_fun_out.loss[0]})
@@ -354,7 +353,8 @@ def finetune(
for ev_batch in tqdm(val_its):
ebatch = process_eval_batch(ev_batch)
ebatch = reshape_batch_for_pmap(ebatch, num_devices)
_, step_fun_out = p_eval_step(replicated_jax_states, eval_prng_seed, ebatch)
_, step_fun_out = p_eval_step(replicated_jax_states, eval_prng_seed,
ebatch)
eval_losses.append(step_fun_out.loss[0])
wandb.log({"eval_step_loss": step_fun_out.loss[0]})
@@ -362,20 +362,17 @@ def finetune(
print(f"Train Loss: {avg_train_loss}, Val Loss: {avg_eval_loss}")
wandb.log(
{
wandb.log({
"epoch": epoch + 1,
"avg_train_loss": avg_train_loss,
"avg_val_loss": avg_eval_loss,
}
)
})
if avg_eval_loss < best_eval_loss or np.isnan(avg_eval_loss):
best_eval_loss = avg_eval_loss
print("Saving checkpoint.")
jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(
replicated_jax_states
)
replicated_jax_states)
if use_lora:
adapter_params = get_adapter_params(
params=jax_state_for_saving.mdl_vars,
@@ -385,9 +382,9 @@ def finetune(
)
jax_state_for_saving.mdl_vars["params"] = adapter_params
checkpoints.save_checkpoint(
jax_state_for_saving, checkpoint_dir, overwrite=True
)
checkpoints.save_checkpoint(jax_state_for_saving,
checkpoint_dir,
overwrite=True)
patience = 0
del jax_state_for_saving
src/adapter/__init__.py
-1
View File
@@ -11,7 +11,6 @@
# 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.
"""adapter init file."""
from .dora_layers import DoraAttentionProjection, DoraCombinedQKVProjection, DoraLinear
src/adapter/dora_layers.py
+7 -8
View File
@@ -21,18 +21,17 @@ WeightHParams = base_layer.WeightHParams
class DoraTheta(base_layer.Theta):
def __init__(self, module):
self.module = module
def _dora_initialized(self):
if (
self.module.has_variable("params", "lora_a")
and self.module.has_variable("params", "lora_b")
and self.module.has_variable("params", "dora_m")
and "lora_a" in self.module._weight_hparams
and "lora_b" in self.module._weight_hparams
and "dora_m" in self.module._weight_hparams
):
if (self.module.has_variable("params", "lora_a") and
self.module.has_variable("params", "lora_b") and
self.module.has_variable("params", "dora_m") and
"lora_a" in self.module._weight_hparams and
"lora_b" in self.module._weight_hparams and
"dora_m" in self.module._weight_hparams):
return True
else:
return False
src/adapter/lora_layers.py
+5 -6
View File
@@ -21,16 +21,15 @@ WeightHParams = base_layer.WeightHParams
class LoraTheta(base_layer.Theta):
def __init__(self, module):
self.module = module
def _lora_initialized(self):
if (
self.module.has_variable("params", "lora_a")
and self.module.has_variable("params", "lora_b")
and "lora_a" in self.module._weight_hparams
and "lora_b" in self.module._weight_hparams
):
if (self.module.has_variable("params", "lora_a") and
self.module.has_variable("params", "lora_b") and
"lora_a" in self.module._weight_hparams and
"lora_b" in self.module._weight_hparams):
return True
else:
return False
src/adapter/utils.py
+51 -81
View File
@@ -11,7 +11,6 @@
# 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.
"""
This file provides functionality for loading and merging adapter weights
in timesfm model, specifically for LoRA and DoRA.
@@ -40,9 +39,10 @@ from adapter.lora_layers import (
from timesfm import TimesFm
def get_adapter_params(
params: dict, lora_target_modules: str, num_layers: int, use_dora: bool = False
) -> dict:
def get_adapter_params(params: dict,
lora_target_modules: str,
num_layers: int,
use_dora: bool = False) -> dict:
"""
Extracts adapter parameters from the given model parameters for saving the checkpoint.
@@ -63,8 +63,7 @@ def get_adapter_params(
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
linear = params["params"]["core_layer"]["stacked_transformer_layer"][
layer_key
]["ff_layer"][ff_layer_key]["linear"]
layer_key]["ff_layer"][ff_layer_key]["linear"]
lora_a = linear["lora_a"]
lora_b = linear["lora_b"]
@@ -79,8 +78,7 @@ def get_adapter_params(
if lora_target_modules in ["all", "attention"]:
attention = params["params"]["core_layer"]["stacked_transformer_layer"][
layer_key
]["self_attention"]
layer_key]["self_attention"]
for component in ["key", "query", "value", "post"]:
lora_a = attention[component]["lora_a"]
@@ -92,9 +90,8 @@ def get_adapter_params(
}
if use_dora:
adapter_params[layer_key][component]["dora_m"] = attention[
component
]["dora_m"]
adapter_params[layer_key][component]["dora_m"] = attention[component][
"dora_m"]
return adapter_params
@@ -118,13 +115,13 @@ def load_adapter_checkpoint(
Returns:
None
"""
"""
currently loading and initializing the model with adapter layers first and then merging the
adapter weights to original weights and replacing the adapter layers back to original layer.
# NOTE: refactor this. there should be a better way to load the LoRA checkpoint.
"""
model._logging(f"Restoring adapter checkpoint from {adapter_checkpoint_path}.")
model._logging(
f"Restoring adapter checkpoint from {adapter_checkpoint_path}.")
start_time = time.time()
original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl = (
load_adapter_layer(
@@ -133,12 +130,10 @@ def load_adapter_checkpoint(
lora_rank=lora_rank,
lora_target_modules=lora_target_modules,
use_dora=use_dora,
)
)
))
var_weight_hparams = model._model.abstract_init_with_metadata(
model._get_sample_inputs(), do_eval=True
)
model._get_sample_inputs(), do_eval=True)
adapter_weight_hparams = _get_adapter_weight_params(
var_weight_hparams=var_weight_hparams,
@@ -179,19 +174,15 @@ def load_adapter_checkpoint(
# replace back with the original model layer
if lora_target_modules in ["all", "mlp"]:
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
original_linear_tpl
)
original_linear_tpl)
if lora_target_modules in ["all", "attention"]:
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
original_attn_tpl
)
original_attn_tpl)
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
original_combined_qkv_tpl
)
original_combined_qkv_tpl)
model._logging(
f"Restored adapter checkpoint in {time.time() - start_time:.2f} seconds."
)
f"Restored adapter checkpoint in {time.time() - start_time:.2f} seconds.")
# jit compile the model
model.jit_decode()
@@ -220,8 +211,8 @@ def _merge_adapter_weights(
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
linear = model._train_state.mdl_vars["params"][
"stacked_transformer_layer"
][layer_key]["ff_layer"][ff_layer_key]["linear"]
"stacked_transformer_layer"][layer_key]["ff_layer"][ff_layer_key][
"linear"]
params = adapter_train_state.mdl_vars[layer_key][ff_layer_key]
lora_a = params["lora_a"]
@@ -249,8 +240,7 @@ def _merge_adapter_weights(
if lora_target_modules in ["all", "attention"]:
attention = model._train_state.mdl_vars["params"][
"stacked_transformer_layer"
][layer_key]["self_attention"]
"stacked_transformer_layer"][layer_key]["self_attention"]
for component in ["key", "query", "value", "post"]:
params = adapter_train_state.mdl_vars[layer_key][component]
@@ -278,9 +268,9 @@ def _merge_adapter_weights(
del attention[component]["lora_b"]
def _get_adapter_weight_params(
var_weight_hparams: dict, lora_target_modules: str, num_layers: int, use_dora: bool
) -> dict:
def _get_adapter_weight_params(var_weight_hparams: dict,
lora_target_modules: str, num_layers: int,
use_dora: bool) -> dict:
"""
Extracts adapter weight parameters from the given variable weight hyperparameters.
@@ -301,8 +291,8 @@ def _get_adapter_weight_params(
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
adapter_weight_params = var_weight_hparams["params"][
"stacked_transformer_layer"
][layer]["ff_layer"][ff_layer_key]["linear"]
"stacked_transformer_layer"][layer]["ff_layer"][ff_layer_key][
"linear"]
adapter_params[layer][ff_layer_key] = {
"lora_a": adapter_weight_params["lora_a"],
"lora_b": adapter_weight_params["lora_b"],
@@ -310,14 +300,12 @@ def _get_adapter_weight_params(
if use_dora:
adapter_params[layer][ff_layer_key]["dora_m"] = (
adapter_weight_params["dora_m"]
)
adapter_weight_params["dora_m"])
if lora_target_modules in ["all", "attention"]:
for component in ["key", "value", "query", "post"]:
adapter_weight_params = var_weight_hparams["params"][
"stacked_transformer_layer"
][layer]["self_attention"][component]
"stacked_transformer_layer"][layer]["self_attention"][component]
adapter_params[layer][component] = {
"lora_a": adapter_weight_params["lora_a"],
"lora_b": adapter_weight_params["lora_b"],
@@ -325,8 +313,7 @@ def _get_adapter_weight_params(
if use_dora:
adapter_params[layer][component]["dora_m"] = adapter_weight_params[
"dora_m"
]
"dora_m"]
return adapter_params
@@ -354,53 +341,41 @@ def load_adapter_layer(
original_linear_tpl = original_attn_tpl = original_combined_qkv_tpl = None
if lora_target_modules in ["all", "mlp"]:
original_linear_tpl = (
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl
)
adapter_linear_tpl = (
pax_fiddle.Config(
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.
tr_fflayer_tpl.fflayer_tpl.linear_tpl)
adapter_linear_tpl = (pax_fiddle.Config(
DoraLinear,
rank=lora_rank,
)
if use_dora
else pax_fiddle.Config(
) if use_dora else pax_fiddle.Config(
LoraLinear,
rank=lora_rank,
)
)
))
adapter_linear_tpl.copy_fields_from(original_linear_tpl)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
adapter_linear_tpl
)
adapter_linear_tpl)
if lora_target_modules in ["all", "attention"]:
original_attn_tpl = (
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl
)
original_attn_tpl = (model.stacked_transformer_params_tpl.
transformer_layer_params_tpl.tr_atten_tpl.proj_tpl)
adapter_attn_tpl = (
pax_fiddle.Config(DoraAttentionProjection, rank=lora_rank)
if use_dora
else pax_fiddle.Config(LoraAttentionProjection, rank=lora_rank)
)
pax_fiddle.Config(DoraAttentionProjection, rank=lora_rank) if use_dora
else pax_fiddle.Config(LoraAttentionProjection, rank=lora_rank))
adapter_attn_tpl.copy_fields_from(original_attn_tpl)
original_combined_qkv_tpl = (
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl
)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.
tr_atten_tpl.combined_qkv_proj_tpl)
adapter_combined_qkv_tpl = (
pax_fiddle.Config(DoraCombinedQKVProjection, rank=lora_rank)
if use_dora
else pax_fiddle.Config(LoraCombinedQKVProjection, rank=lora_rank)
)
pax_fiddle.Config(DoraCombinedQKVProjection, rank=lora_rank) if use_dora
else pax_fiddle.Config(LoraCombinedQKVProjection, rank=lora_rank))
adapter_combined_qkv_tpl.copy_fields_from(original_combined_qkv_tpl)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
adapter_attn_tpl
)
adapter_attn_tpl)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
adapter_combined_qkv_tpl
)
adapter_combined_qkv_tpl)
# initialize and add adapter weights
_initialize_adapter_params(
@@ -441,16 +416,14 @@ def _initialize_adapter_params(
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
linear = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
"ff_layer"
][ff_layer_key]["linear"]
"ff_layer"][ff_layer_key]["linear"]
original_w = linear["w"]
input_dim, output_dim = original_w.shape
std_dev = 1 / jnp.sqrt(lora_rank)
normal_initializer = jax.nn.initializers.normal(std_dev)
lora_a = normal_initializer(
jax.random.key(seed), (input_dim, lora_rank), jnp.float32
)
lora_a = normal_initializer(jax.random.key(seed),
(input_dim, lora_rank), jnp.float32)
lora_b = jnp.zeros((output_dim, lora_rank))
linear["lora_a"] = lora_a
@@ -462,8 +435,7 @@ def _initialize_adapter_params(
if lora_target_modules in ["all", "attention"]:
attention = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
"self_attention"
]
"self_attention"]
for component in ["key", "query", "value", "post"]:
original_w = attention[component]["w"]
@@ -471,17 +443,15 @@ def _initialize_adapter_params(
std_dev = 1 / jnp.sqrt(lora_rank)
normal_initializer = jax.nn.initializers.normal(std_dev)
lora_a = normal_initializer(
jax.random.key(seed), (w_dim, lora_rank), jnp.float32
)
lora_a = normal_initializer(jax.random.key(seed), (w_dim, lora_rank),
jnp.float32)
lora_b = jnp.zeros((w_dim, lora_rank))
attention[component]["lora_a"] = lora_a
attention[component]["lora_b"] = lora_b
if use_dora:
norm = jnp.linalg.norm(
original_w, ord=2, axis=0, keepdims=True
).astype(jnp.float32)
norm = jnp.linalg.norm(original_w, ord=2, axis=0,
keepdims=True).astype(jnp.float32)
attention[component]["dora_m"] = norm
return mdl_vars
src/finetuning/finetuning_example.py
+4 -3
View File
@@ -43,11 +43,11 @@ flags.DEFINE_list(
)
flags.DEFINE_string(
"local_model_path",
None,
"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."""
@@ -161,7 +161,8 @@ def get_model(load_weights: bool = False):
tfm_config = tfm._model_config
model = PatchedTimeSeriesDecoder(tfm_config)
checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
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)
src/timesfm/__init__.py
+3 -1
View File
@@ -32,4 +32,6 @@ try:
except Exception as _:
from timesfm.timesfm_torch import TimesFmTorch as TimesFm
print(f"Loaded PyTorch TimesFM, likely because python version is {sys.version}.")
print(
f"Loaded PyTorch TimesFM, likely because python version is {sys.version}."
)
src/timesfm/time_features.py
+13 -13
View File
@@ -11,7 +11,6 @@
# 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.
"""Directory to extract time covariates.
Extract time covariates from datetime.
@@ -36,7 +35,6 @@ from pandas.tseries.offsets import Easter
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm
# This is 183 to cover half a year (in both directions), also for leap years
# + 17 as Eastern can be between March, 22 - April, 25
MAX_WINDOW = 183 + 17
@@ -50,8 +48,7 @@ def _distance_to_holiday(holiday):
index - pd.Timedelta(days=MAX_WINDOW),
index + pd.Timedelta(days=MAX_WINDOW),
)
assert (
len(holiday_date) != 0 # pylint: disable=g-explicit-length-test
assert (len(holiday_date) != 0 # pylint: disable=g-explicit-length-test
), f"No closest holiday for the date index {index} found."
# It sometimes returns two dates if it is exactly half a year after the
# holiday. In this case, the smaller distance (182 days) is returned.
@@ -60,16 +57,19 @@ def _distance_to_holiday(holiday):
return _distance_to_day
EasterSunday = Holiday(
"Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
)
EasterSunday = Holiday("Easter Sunday",
month=1,
day=1,
offset=[Easter(), Day(0)])
NewYearsDay = Holiday("New Years Day", month=1, day=1)
SuperBowl = Holiday(
"Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1))
)
MothersDay = Holiday(
"Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2))
)
SuperBowl = Holiday("Superbowl",
month=2,
day=1,
offset=DateOffset(weekday=SU(1)))
MothersDay = Holiday("Mothers Day",
month=5,
day=1,
offset=DateOffset(weekday=SU(2)))
IndependenceDay = Holiday("Independence Day", month=7, day=4)
ChristmasEve = Holiday("Christmas", month=12, day=24)
ChristmasDay = Holiday("Christmas", month=12, day=25)
src/timesfm/timesfm_base.py
+4 -11
View File
@@ -57,18 +57,11 @@ def freq_map(freq: str):
return 1
elif freq.endswith(("H", "T", "MIN", "D", "B", "U", "S")):
return 0
elif (
freq.endswith(("W", "M"))
or freq.startswith("W-")
or (freq.startswith("M") and len(freq) == 2)
):
elif (freq.endswith(("W", "M")) or freq.startswith("W-") or
(freq.startswith("M") and len(freq) == 2)):
return 1
elif (
freq.endswith(("Y", "Q", "A"))
or freq.startswith("Y-")
or freq.startswith("Q-")
or freq.startswith("A-")
):
elif (freq.endswith(("Y", "Q", "A")) or freq.startswith("Y-") or
freq.startswith("Q-") or freq.startswith("A-")):
return 2
else:
raise ValueError(f"Invalid frequency: {freq}")
tests/test_timesfm.py
+4 -6
View File
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from datetime import datetime, timedelta
import numpy as np
@@ -22,9 +21,9 @@ import pytest
import timesfm
def create_sample_dataframe(
start_date: datetime, end_date: datetime, freq: str = "D"
) -> pd.DataFrame:
def create_sample_dataframe(start_date: datetime,
end_date: datetime,
freq: str = "D") -> pd.DataFrame:
"""
Create a sample DataFrame with time series data.
@@ -75,8 +74,7 @@ def test_timesfm_forecast_on_df(
assert (
len(forecast_df) == prediction_length
), f"Expected forecast length of {prediction_length}, but got {len(forecast_df)}"
assert (
"timesfm" in forecast_df.columns
assert ("timesfm" in forecast_df.columns
), "Forecast DataFrame should contain 'timesfm' column"
last_input_date = input_df["ds"].max()