Merge pull request #104 from tanmayshishodia/feature/lora
LoRA and DoRA PEFT support for Fine-Tuning TimesFM
This commit is contained in:
Rajat Sen
+6
-1
@@ -1,3 +1,8 @@
|
||||
.venv/
|
||||
dist/
|
||||
**__pycache__/** */
|
||||
__pycache__/
|
||||
checkpoints/
|
||||
wandb/
|
||||
datasets/
|
||||
results/
|
||||
timesfm_jax.egg-info/
|
||||
|
||||
@@ -16,3 +16,6 @@ dependencies:
|
||||
- jax[cuda12]==0.4.26
|
||||
- einshape
|
||||
- scikit-learn
|
||||
- typer
|
||||
- wandb
|
||||
- pytest
|
||||
|
||||
@@ -16,3 +16,6 @@ dependencies:
|
||||
- jax[cpu]==0.4.26
|
||||
- einshape
|
||||
- scikit-learn
|
||||
- typer
|
||||
- wandb
|
||||
- pytest
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
# Fine-Tuning Pipeline
|
||||
|
||||
This folder contains a generic fine-tuning pipeline designed to support multiple PEFT fine-tuning strategies.
|
||||
|
||||
## Features
|
||||
|
||||
- **Supported Fine-Tuning Strategies**:
|
||||
- **Full Fine-Tuning**: Adjusts all parameters of the model during training.
|
||||
- **[Linear Probing](https://arxiv.org/abs/2302.11939)**: Fine-tunes only the residual blocks and the embedding layer, leaving other parameters unchanged.
|
||||
- **[LoRA (Low-Rank Adaptation)](https://arxiv.org/abs/2106.09685)**: A memory-efficient method that fine-tunes a small number of parameters by decomposing the weight matrices into low-rank matrices.
|
||||
- **[DoRA (Directional LoRA)](https://arxiv.org/abs/2402.09353v4)**: An extension of LoRA that decomposes pre-trained weights into magnitude and direction components. It uses LoRA for directional adaptation, enhancing learning capacity and stability without additional inference overhead.
|
||||
|
||||
## Usage
|
||||
### Fine-Tuning Script
|
||||
The provided finetune.py script allows you to fine-tune a model with specific configurations. You can customize various parameters to suit your dataset and desired fine-tuning strategy.
|
||||
|
||||
Example Usage:
|
||||
|
||||
```zsh
|
||||
source finetune.sh
|
||||
```
|
||||
This script runs the finetune.py file with a predefined set of hyperparameters for the model. You can adjust the parameters in the script as needed.
|
||||
|
||||
### Available Options
|
||||
Run the script with the --help flag to see a full list of available options and their descriptions:
|
||||
```zsh
|
||||
python3 finetune.py --help
|
||||
```
|
||||
Script Configuration
|
||||
You can modify the following key parameters directly in the finetune.sh script:
|
||||
Fine-Tuning Strategy: Toggle between full fine-tuning, LoRA \[`--use-lora`\], DoRA [\[`--use-dora`\]], or Linear Probing \[`--use-linear-probing`\].
|
||||
|
||||
### Performance Comparison
|
||||
The figure below compares the performance of LoRA/DoRA against Linear Probing under the following conditions:
|
||||
|
||||
<img width="528" alt="image" src="https://github.com/user-attachments/assets/6c9f820b-5865-4821-8014-c346b9d632a5">
|
||||
|
||||
- Training data split: 60% train, 20% validation, 20% test.
|
||||
- Benchmark: context_len=128, horizon_len=96
|
||||
- Fine-tuning: context_len=128, horizon_len=128
|
||||
- Black: Best result.
|
||||
- Blue: Second best result.
|
||||
@@ -0,0 +1,402 @@
|
||||
# Copyright 2024 The Google Research Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Finetune pipeline.
|
||||
"""
|
||||
import gc
|
||||
import logging
|
||||
import warnings
|
||||
from datetime import datetime
|
||||
from typing import Tuple
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import typer
|
||||
import wandb
|
||||
from jax import numpy as jnp
|
||||
from paxml import checkpoint_types, checkpoints, learners, tasks_lib, trainer_lib
|
||||
from praxis import optimizers, pax_fiddle, py_utils, schedules
|
||||
from rich import print
|
||||
from tqdm import tqdm
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from adapter.utils import get_adapter_params, load_adapter_layer
|
||||
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.
|
||||
"""
|
||||
INPUT_PATCH_LEN = 32
|
||||
OUTPUT_PATCH_LEN = 128
|
||||
NUM_LAYERS = 20
|
||||
MODEL_DIMS = 1280
|
||||
|
||||
QUANTILES = list(np.arange(1, 10) / 10.0)
|
||||
EPS = 1e-7
|
||||
RANDOM_SEED = 1234
|
||||
|
||||
|
||||
def finetune(
|
||||
*,
|
||||
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")],
|
||||
horizon_len: Annotated[int, typer.Option(help="Prediction length.")],
|
||||
freq: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
...,
|
||||
help="Frequency Map Str",
|
||||
),
|
||||
],
|
||||
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",
|
||||
),
|
||||
] = (0, 0, 0),
|
||||
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,
|
||||
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")],
|
||||
early_stop_patience: Annotated[
|
||||
int, typer.Option(..., help="Early stopping patience")
|
||||
] = 5,
|
||||
use_lora: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
help="Train low rank adapters for stacked transformer block",
|
||||
),
|
||||
] = False,
|
||||
lora_rank: Annotated[
|
||||
int,
|
||||
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]"
|
||||
),
|
||||
] = "all",
|
||||
use_dora: Annotated[
|
||||
bool,
|
||||
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.",
|
||||
),
|
||||
] = 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",
|
||||
) -> None:
|
||||
key = jax.random.PRNGKey(seed=RANDOM_SEED)
|
||||
wandb.init(project=wandb_project, config=locals())
|
||||
|
||||
data_df = pd.read_csv(open(data_path, "r"))
|
||||
|
||||
if boundaries == (0, 0, 0):
|
||||
# Default boundaries: train 60%, val 20%, test 20%
|
||||
boundaries = [
|
||||
int(len(data_df) * 0.6),
|
||||
int(len(data_df) * 0.8),
|
||||
len(data_df) - 1,
|
||||
]
|
||||
|
||||
ts_cols = [col for col in data_df.columns if col != datetime_col]
|
||||
|
||||
dtl = data_loader.TimeSeriesdata(
|
||||
data_path=data_path,
|
||||
datetime_col=datetime_col,
|
||||
num_cov_cols=None,
|
||||
cat_cov_cols=None,
|
||||
ts_cols=np.array(ts_cols),
|
||||
train_range=[0, boundaries[0]],
|
||||
val_range=[boundaries[0], boundaries[1]],
|
||||
test_range=[boundaries[1], boundaries[2]],
|
||||
hist_len=context_len,
|
||||
pred_len=horizon_len,
|
||||
batch_size=batch_size,
|
||||
freq=freq,
|
||||
normalize=normalize,
|
||||
epoch_len=None,
|
||||
holiday=False,
|
||||
permute=False,
|
||||
)
|
||||
|
||||
train_batches = dtl.tf_dataset(mode="train", shift=1).batch(batch_size)
|
||||
val_batches = dtl.tf_dataset(mode="val", shift=horizon_len)
|
||||
|
||||
for tbatch in tqdm(train_batches.as_numpy_iterator()):
|
||||
pass
|
||||
|
||||
tfm = TimesFm(
|
||||
context_len=context_len,
|
||||
horizon_len=horizon_len,
|
||||
input_patch_len=INPUT_PATCH_LEN,
|
||||
output_patch_len=OUTPUT_PATCH_LEN,
|
||||
num_layers=NUM_LAYERS,
|
||||
model_dims=MODEL_DIMS,
|
||||
backend=backend,
|
||||
per_core_batch_size=batch_size,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
|
||||
if checkpoint_path:
|
||||
tfm.load_from_checkpoint(
|
||||
checkpoint_path=checkpoint_path,
|
||||
checkpoint_type=checkpoints.CheckpointType.FLAX,
|
||||
)
|
||||
else:
|
||||
tfm.load_from_checkpoint(
|
||||
repo_id=model_name,
|
||||
checkpoint_type=checkpoints.CheckpointType.FLAX,
|
||||
)
|
||||
|
||||
model = pax_fiddle.Config(
|
||||
patched_decoder.PatchedDecoderFinetuneModel,
|
||||
name="patched_decoder_finetune",
|
||||
core_layer_tpl=tfm.model_p,
|
||||
)
|
||||
|
||||
if use_lora:
|
||||
load_adapter_layer(
|
||||
mdl_vars=tfm._train_state.mdl_vars,
|
||||
model=model.core_layer_tpl,
|
||||
lora_rank=lora_rank,
|
||||
lora_target_modules=lora_target_modules,
|
||||
use_dora=use_dora,
|
||||
)
|
||||
|
||||
@pax_fiddle.auto_config
|
||||
def build_learner() -> learners.Learner:
|
||||
bprop_variable_inclusion = []
|
||||
bprop_variable_exclusion = []
|
||||
if use_lora:
|
||||
bprop_variable_inclusion.append(r"^.*lora.*$")
|
||||
if use_dora:
|
||||
bprop_variable_inclusion.append(r"^.*dora.*$")
|
||||
elif use_linear_probing:
|
||||
bprop_variable_exclusion = [".*/stacked_transformer_layer/.*"]
|
||||
|
||||
return pax_fiddle.Config(
|
||||
learners.Learner,
|
||||
name="learner",
|
||||
loss_name="avg_qloss",
|
||||
optimizer=optimizers.Adam(
|
||||
epsilon=adam_epsilon,
|
||||
clip_threshold=adam_clip_threshold,
|
||||
learning_rate=learning_rate,
|
||||
lr_schedule=pax_fiddle.Config(
|
||||
schedules.Cosine,
|
||||
initial_value=cos_initial_decay_value,
|
||||
final_value=cos_final_decay_value,
|
||||
total_steps=cos_decay_steps,
|
||||
),
|
||||
ema_decay=ema_decay,
|
||||
),
|
||||
bprop_variable_exclusion=bprop_variable_exclusion,
|
||||
bprop_variable_inclusion=bprop_variable_inclusion,
|
||||
)
|
||||
|
||||
task_p = tasks_lib.SingleTask(
|
||||
name="ts-learn",
|
||||
model=model,
|
||||
train=tasks_lib.SingleTask.Train(
|
||||
learner=build_learner(),
|
||||
),
|
||||
)
|
||||
|
||||
task_p.model.ici_mesh_shape = [1, 1, 1]
|
||||
task_p.model.mesh_axis_names = ["replica", "data", "mdl"]
|
||||
|
||||
DEVICES = np.array(jax.devices()).reshape([1, 1, 1])
|
||||
jax.sharding.Mesh(DEVICES, ["replica", "data", "mdl"])
|
||||
|
||||
num_devices = jax.local_device_count()
|
||||
print(f"num_devices: {num_devices}")
|
||||
print(f"device kind: {jax.local_devices()[0].device_kind}")
|
||||
|
||||
jax_task = task_p
|
||||
key, init_key = jax.random.split(key)
|
||||
|
||||
def process_train_batch(batch):
|
||||
past_ts = batch[0].reshape(batch_size * len(ts_cols), -1)
|
||||
actual_ts = batch[3].reshape(batch_size * len(ts_cols), -1)
|
||||
return NestedMap(input_ts=past_ts, actual_ts=actual_ts)
|
||||
|
||||
def process_eval_batch(batch):
|
||||
past_ts = batch[0]
|
||||
actual_ts = batch[3]
|
||||
return NestedMap(input_ts=past_ts, actual_ts=actual_ts)
|
||||
|
||||
jax_model_states, _ = trainer_lib.initialize_model_state(
|
||||
jax_task,
|
||||
init_key,
|
||||
process_train_batch(tbatch),
|
||||
checkpoint_type=checkpoint_types.CheckpointType.GDA,
|
||||
)
|
||||
jax_model_states.mdl_vars["params"]["core_layer"] = tfm._train_state.mdl_vars[
|
||||
"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)
|
||||
|
||||
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)
|
||||
|
||||
key, train_key, eval_key = jax.random.split(key, 3)
|
||||
train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())
|
||||
eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())
|
||||
|
||||
p_train_step = jax.pmap(train_step, axis_name="batch")
|
||||
p_eval_step = jax.pmap(eval_step, axis_name="batch")
|
||||
|
||||
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:])
|
||||
nbsize = bsize // num_devices
|
||||
return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)
|
||||
|
||||
return jax.tree.map(_reshape, batch)
|
||||
|
||||
patience = 0
|
||||
best_eval_loss = 1e7
|
||||
checkpoint_dir = f"{checkpoint_dir}/run_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{wandb.run.id}"
|
||||
for epoch in range(num_epochs):
|
||||
if patience >= early_stop_patience:
|
||||
print("Early stopping.")
|
||||
break
|
||||
print(f"Epoch: {epoch + 1}")
|
||||
train_its = train_batches.as_numpy_iterator()
|
||||
train_losses = []
|
||||
for batch in tqdm(train_its):
|
||||
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
|
||||
)
|
||||
train_losses.append(step_fun_out.loss[0])
|
||||
wandb.log({"train_step_loss": step_fun_out.loss[0]})
|
||||
|
||||
avg_train_loss = np.mean(train_losses)
|
||||
|
||||
print("Starting eval.")
|
||||
val_its = val_batches.as_numpy_iterator()
|
||||
eval_losses = []
|
||||
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)
|
||||
eval_losses.append(step_fun_out.loss[0])
|
||||
wandb.log({"eval_step_loss": step_fun_out.loss[0]})
|
||||
|
||||
avg_eval_loss = np.mean(eval_losses)
|
||||
|
||||
print(f"Train Loss: {avg_train_loss}, Val Loss: {avg_eval_loss}")
|
||||
|
||||
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
|
||||
)
|
||||
if use_lora:
|
||||
adapter_params = get_adapter_params(
|
||||
params=jax_state_for_saving.mdl_vars,
|
||||
lora_target_modules=lora_target_modules,
|
||||
num_layers=NUM_LAYERS,
|
||||
use_dora=use_dora,
|
||||
)
|
||||
jax_state_for_saving.mdl_vars["params"] = adapter_params
|
||||
|
||||
checkpoints.save_checkpoint(
|
||||
jax_state_for_saving, checkpoint_dir, overwrite=True
|
||||
)
|
||||
|
||||
patience = 0
|
||||
del jax_state_for_saving
|
||||
gc.collect()
|
||||
else:
|
||||
patience += 1
|
||||
print(f"patience: {patience}")
|
||||
print("Fine-tuning completed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
typer.run(finetune)
|
||||
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Script to finetune a model with specific configurations
|
||||
# Adjust the parameters below as needed. For a full list of options and descriptions, run the script with the --help flag.
|
||||
|
||||
export TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false
|
||||
|
||||
python3 finetune.py \
|
||||
--model-name="google/timesfm-1.0-200m" \
|
||||
--backend="cpu" \
|
||||
--horizon-len=128 \
|
||||
--context-len=512 \
|
||||
--freq="15min" \
|
||||
--data-path="../datasets/ETT-small/ETTm1.csv" \
|
||||
--num-epochs=100 \
|
||||
--learning-rate=1e-3 \
|
||||
--adam-epsilon=1e-7 \
|
||||
--adam-clip-threshold=1e2 \
|
||||
--early-stop-patience=10 \
|
||||
--datetime-col="date" \
|
||||
--use-lora \
|
||||
--lora-rank=1 \
|
||||
--lora-target-modules="all" \
|
||||
--use-dora \
|
||||
--cos-initial-decay-value=1e-4 \
|
||||
--cos-decay-steps=40000 \
|
||||
--cos-final-decay-value=1e-5 \
|
||||
--ema-decay=0.9999
|
||||
|
||||
# To see all available options and their descriptions, use the --help flag
|
||||
# python3 finetune.py --help
|
||||
@@ -0,0 +1,203 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Load Base Model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from timesfm import TimesFm, freq_map, data_loader\n",
|
||||
"from adapter.utils import load_adapter_checkpoint\n",
|
||||
"from tqdm import tqdm\n",
|
||||
"import numpy as np\n",
|
||||
"import pandas as pd\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"tfm = TimesFm(\n",
|
||||
" context_len=512,\n",
|
||||
" horizon_len=128,\n",
|
||||
" input_patch_len=32,\n",
|
||||
" output_patch_len=128,\n",
|
||||
" num_layers=20,\n",
|
||||
" model_dims=1280,\n",
|
||||
" backend=\"cpu\",\n",
|
||||
")\n",
|
||||
"tfm.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"DATA_DICT = {\n",
|
||||
" \"ettm2\": {\n",
|
||||
" \"boundaries\": [34560, 46080, 57600],\n",
|
||||
" \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
|
||||
" \"freq\": \"15min\",\n",
|
||||
" },\n",
|
||||
" \"ettm1\": {\n",
|
||||
" \"boundaries\": [34560, 46080, 57600],\n",
|
||||
" \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
|
||||
" \"freq\": \"15min\",\n",
|
||||
" },\n",
|
||||
" \"etth2\": {\n",
|
||||
" \"boundaries\": [8640, 11520, 14400],\n",
|
||||
" \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
|
||||
" \"freq\": \"H\",\n",
|
||||
" },\n",
|
||||
" \"etth1\": {\n",
|
||||
" \"boundaries\": [8640, 11520, 14400],\n",
|
||||
" \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
|
||||
" \"freq\": \"H\",\n",
|
||||
" },\n",
|
||||
" \"elec\": {\n",
|
||||
" \"boundaries\": [18413, 21044, 26304],\n",
|
||||
" \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
|
||||
" \"freq\": \"H\",\n",
|
||||
" },\n",
|
||||
" \"traffic\": {\n",
|
||||
" \"boundaries\": [12280, 14036, 17544],\n",
|
||||
" \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
|
||||
" \"freq\": \"H\",\n",
|
||||
" },\n",
|
||||
" \"weather\": {\n",
|
||||
" \"boundaries\": [36887, 42157, 52696],\n",
|
||||
" \"data_path\": \"../datasets/weather/weather.csv\",\n",
|
||||
" \"freq\": \"10min\",\n",
|
||||
" },\n",
|
||||
"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Load Adapter Checkpoint\n",
|
||||
"\n",
|
||||
"Specify the adapter checkpoint path, rank and the modules used to train the adapters and whether dora was employed or not."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"load_adapter_checkpoint(\n",
|
||||
" model=tfm,\n",
|
||||
" adapter_checkpoint_path=\"./checkpoints/run_20240716_163900_lyo4psz3\",\n",
|
||||
" lora_rank=1,\n",
|
||||
" lora_target_modules=\"all\",\n",
|
||||
" use_dora=True,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Test Performance"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"dataset = \"ettm1\"\n",
|
||||
"data_path = DATA_DICT[dataset][\"data_path\"]\n",
|
||||
"freq = DATA_DICT[dataset][\"freq\"]\n",
|
||||
"int_freq = freq_map(freq)\n",
|
||||
"boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
|
||||
"\n",
|
||||
"data_df = pd.read_csv(open(data_path, \"r\"))\n",
|
||||
"\n",
|
||||
"ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
|
||||
"num_cov_cols = None\n",
|
||||
"cat_cov_cols = None\n",
|
||||
"\n",
|
||||
"context_len = 512\n",
|
||||
"pred_len = 96\n",
|
||||
"\n",
|
||||
"num_ts = len(ts_cols)\n",
|
||||
"batch_size = 16\n",
|
||||
"\n",
|
||||
"dtl = data_loader.TimeSeriesdata(\n",
|
||||
" data_path=data_path,\n",
|
||||
" datetime_col=\"date\",\n",
|
||||
" num_cov_cols=num_cov_cols,\n",
|
||||
" cat_cov_cols=cat_cov_cols,\n",
|
||||
" ts_cols=np.array(ts_cols),\n",
|
||||
" train_range=[0, boundaries[0]],\n",
|
||||
" val_range=[boundaries[0], boundaries[1]],\n",
|
||||
" test_range=[boundaries[1], boundaries[2]],\n",
|
||||
" hist_len=context_len,\n",
|
||||
" pred_len=pred_len,\n",
|
||||
" batch_size=num_ts,\n",
|
||||
" freq=\"15min\",\n",
|
||||
" normalize=True,\n",
|
||||
" epoch_len=None,\n",
|
||||
" holiday=False,\n",
|
||||
" permute=True,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"mae_losses = []\n",
|
||||
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
|
||||
" past = batch[0]\n",
|
||||
" actuals = batch[3]\n",
|
||||
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
|
||||
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
|
||||
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
|
||||
"\n",
|
||||
"print(f\"MAE: {np.mean(mae_losses)}\")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "tanmay_tfm_env",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.14"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Generated
+326
-1
@@ -732,6 +732,20 @@ files = [
|
||||
{file = "dm_tree-0.1.8-cp39-cp39-win_amd64.whl", hash = "sha256:8ed3564abed97c806db122c2d3e1a2b64c74a63debe9903aad795167cc301368"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "docker-pycreds"
|
||||
version = "0.4.0"
|
||||
description = "Python bindings for the docker credentials store API"
|
||||
optional = false
|
||||
python-versions = "*"
|
||||
files = [
|
||||
{file = "docker-pycreds-0.4.0.tar.gz", hash = "sha256:6ce3270bcaf404cc4c3e27e4b6c70d3521deae82fb508767870fdbf772d584d4"},
|
||||
{file = "docker_pycreds-0.4.0-py2.py3-none-any.whl", hash = "sha256:7266112468627868005106ec19cd0d722702d2b7d5912a28e19b826c3d37af49"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
six = ">=1.4.0"
|
||||
|
||||
[[package]]
|
||||
name = "docstring-parser"
|
||||
version = "0.16"
|
||||
@@ -1165,6 +1179,38 @@ testing = ["absl-py (>=0.1.6)", "mock (>=3.0.5)", "nose"]
|
||||
tf-nightly = ["tf-nightly"]
|
||||
torch = ["torch (>=1.3.0)"]
|
||||
|
||||
[[package]]
|
||||
name = "gitdb"
|
||||
version = "4.0.11"
|
||||
description = "Git Object Database"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "gitdb-4.0.11-py3-none-any.whl", hash = "sha256:81a3407ddd2ee8df444cbacea00e2d038e40150acfa3001696fe0dcf1d3adfa4"},
|
||||
{file = "gitdb-4.0.11.tar.gz", hash = "sha256:bf5421126136d6d0af55bc1e7c1af1c397a34f5b7bd79e776cd3e89785c2b04b"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
smmap = ">=3.0.1,<6"
|
||||
|
||||
[[package]]
|
||||
name = "gitpython"
|
||||
version = "3.1.43"
|
||||
description = "GitPython is a Python library used to interact with Git repositories"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "GitPython-3.1.43-py3-none-any.whl", hash = "sha256:eec7ec56b92aad751f9912a73404bc02ba212a23adb2c7098ee668417051a1ff"},
|
||||
{file = "GitPython-3.1.43.tar.gz", hash = "sha256:35f314a9f878467f5453cc1fee295c3e18e52f1b99f10f6cf5b1682e968a9e7c"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
gitdb = ">=4.0.1,<5"
|
||||
|
||||
[package.extras]
|
||||
doc = ["sphinx (==4.3.2)", "sphinx-autodoc-typehints", "sphinx-rtd-theme", "sphinxcontrib-applehelp (>=1.0.2,<=1.0.4)", "sphinxcontrib-devhelp (==1.0.2)", "sphinxcontrib-htmlhelp (>=2.0.0,<=2.0.1)", "sphinxcontrib-qthelp (==1.0.3)", "sphinxcontrib-serializinghtml (==1.1.5)"]
|
||||
test = ["coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock", "mypy", "pre-commit", "pytest (>=7.3.1)", "pytest-cov", "pytest-instafail", "pytest-mock", "pytest-sugar", "typing-extensions"]
|
||||
|
||||
[[package]]
|
||||
name = "google-auth"
|
||||
version = "2.32.0"
|
||||
@@ -1484,6 +1530,17 @@ files = [
|
||||
docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (<7.2.5)", "sphinx (>=3.5)", "sphinx-lint"]
|
||||
testing = ["jaraco.test (>=5.4)", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-mypy", "pytest-ruff (>=0.2.1)", "zipp (>=3.17)"]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.0.0"
|
||||
description = "brain-dead simple config-ini parsing"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"},
|
||||
{file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "inquirerpy"
|
||||
version = "0.3.4"
|
||||
@@ -3613,6 +3670,21 @@ docs = ["furo (>=2023.9.10)", "proselint (>=0.13)", "sphinx (>=7.2.6)", "sphinx-
|
||||
test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.4.3)", "pytest-cov (>=4.1)", "pytest-mock (>=3.12)"]
|
||||
type = ["mypy (>=1.8)"]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.5.0"
|
||||
description = "plugin and hook calling mechanisms for python"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "pluggy-1.5.0-py3-none-any.whl", hash = "sha256:44e1ad92c8ca002de6377e165f3e0f1be63266ab4d554740532335b9d75ea669"},
|
||||
{file = "pluggy-1.5.0.tar.gz", hash = "sha256:2cffa88e94fdc978c4c574f15f9e59b7f4201d439195c3715ca9e2486f1d0cf1"},
|
||||
]
|
||||
|
||||
[package.extras]
|
||||
dev = ["pre-commit", "tox"]
|
||||
testing = ["pytest", "pytest-benchmark"]
|
||||
|
||||
[[package]]
|
||||
name = "portalocker"
|
||||
version = "2.10.1"
|
||||
@@ -3876,6 +3948,28 @@ files = [
|
||||
[package.extras]
|
||||
diagrams = ["jinja2", "railroad-diagrams"]
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "8.3.2"
|
||||
description = "pytest: simple powerful testing with Python"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "pytest-8.3.2-py3-none-any.whl", hash = "sha256:4ba08f9ae7dcf84ded419494d229b48d0903ea6407b030eaec46df5e6a73bba5"},
|
||||
{file = "pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
colorama = {version = "*", markers = "sys_platform == \"win32\""}
|
||||
exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""}
|
||||
iniconfig = "*"
|
||||
packaging = "*"
|
||||
pluggy = ">=1.5,<2"
|
||||
tomli = {version = ">=1", markers = "python_version < \"3.11\""}
|
||||
|
||||
[package.extras]
|
||||
dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
@@ -4777,6 +4871,56 @@ files = [
|
||||
{file = "sentencepiece-0.1.99.tar.gz", hash = "sha256:189c48f5cb2949288f97ccdb97f0473098d9c3dcf5a3d99d4eabe719ec27297f"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sentry-sdk"
|
||||
version = "2.12.0"
|
||||
description = "Python client for Sentry (https://sentry.io)"
|
||||
optional = false
|
||||
python-versions = ">=3.6"
|
||||
files = [
|
||||
{file = "sentry_sdk-2.12.0-py2.py3-none-any.whl", hash = "sha256:7a8d5163d2ba5c5f4464628c6b68f85e86972f7c636acc78aed45c61b98b7a5e"},
|
||||
{file = "sentry_sdk-2.12.0.tar.gz", hash = "sha256:8763840497b817d44c49b3fe3f5f7388d083f2337ffedf008b2cdb63b5c86dc6"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
certifi = "*"
|
||||
urllib3 = ">=1.26.11"
|
||||
|
||||
[package.extras]
|
||||
aiohttp = ["aiohttp (>=3.5)"]
|
||||
anthropic = ["anthropic (>=0.16)"]
|
||||
arq = ["arq (>=0.23)"]
|
||||
asyncpg = ["asyncpg (>=0.23)"]
|
||||
beam = ["apache-beam (>=2.12)"]
|
||||
bottle = ["bottle (>=0.12.13)"]
|
||||
celery = ["celery (>=3)"]
|
||||
celery-redbeat = ["celery-redbeat (>=2)"]
|
||||
chalice = ["chalice (>=1.16.0)"]
|
||||
clickhouse-driver = ["clickhouse-driver (>=0.2.0)"]
|
||||
django = ["django (>=1.8)"]
|
||||
falcon = ["falcon (>=1.4)"]
|
||||
fastapi = ["fastapi (>=0.79.0)"]
|
||||
flask = ["blinker (>=1.1)", "flask (>=0.11)", "markupsafe"]
|
||||
grpcio = ["grpcio (>=1.21.1)", "protobuf (>=3.8.0)"]
|
||||
httpx = ["httpx (>=0.16.0)"]
|
||||
huey = ["huey (>=2)"]
|
||||
huggingface-hub = ["huggingface-hub (>=0.22)"]
|
||||
langchain = ["langchain (>=0.0.210)"]
|
||||
loguru = ["loguru (>=0.5)"]
|
||||
openai = ["openai (>=1.0.0)", "tiktoken (>=0.3.0)"]
|
||||
opentelemetry = ["opentelemetry-distro (>=0.35b0)"]
|
||||
opentelemetry-experimental = ["opentelemetry-distro"]
|
||||
pure-eval = ["asttokens", "executing", "pure-eval"]
|
||||
pymongo = ["pymongo (>=3.1)"]
|
||||
pyspark = ["pyspark (>=2.4.4)"]
|
||||
quart = ["blinker (>=1.1)", "quart (>=0.16.1)"]
|
||||
rq = ["rq (>=0.6)"]
|
||||
sanic = ["sanic (>=0.8)"]
|
||||
sqlalchemy = ["sqlalchemy (>=1.2)"]
|
||||
starlette = ["starlette (>=0.19.1)"]
|
||||
starlite = ["starlite (>=1.48)"]
|
||||
tornado = ["tornado (>=6)"]
|
||||
|
||||
[[package]]
|
||||
name = "seqio-nightly"
|
||||
version = "0.0.17.dev20231010"
|
||||
@@ -4807,6 +4951,106 @@ cache-tasks = ["apache-beam"]
|
||||
gcp = ["gevent", "google-api-python-client", "google-cloud-storage", "google-compute-engine", "oauth2client"]
|
||||
test = ["pytest"]
|
||||
|
||||
[[package]]
|
||||
name = "setproctitle"
|
||||
version = "1.3.3"
|
||||
description = "A Python module to customize the process title"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:897a73208da48db41e687225f355ce993167079eda1260ba5e13c4e53be7f754"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8c331e91a14ba4076f88c29c777ad6b58639530ed5b24b5564b5ed2fd7a95452"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bbbd6c7de0771c84b4aa30e70b409565eb1fc13627a723ca6be774ed6b9d9fa3"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c05ac48ef16ee013b8a326c63e4610e2430dbec037ec5c5b58fcced550382b74"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1342f4fdb37f89d3e3c1c0a59d6ddbedbde838fff5c51178a7982993d238fe4f"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc74e84fdfa96821580fb5e9c0b0777c1c4779434ce16d3d62a9c4d8c710df39"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9617b676b95adb412bb69645d5b077d664b6882bb0d37bfdafbbb1b999568d85"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:6a249415f5bb88b5e9e8c4db47f609e0bf0e20a75e8d744ea787f3092ba1f2d0"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:38da436a0aaace9add67b999eb6abe4b84397edf4a78ec28f264e5b4c9d53cd5"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:da0d57edd4c95bf221b2ebbaa061e65b1788f1544977288bdf95831b6e44e44d"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-win32.whl", hash = "sha256:a1fcac43918b836ace25f69b1dca8c9395253ad8152b625064415b1d2f9be4fb"},
|
||||
{file = "setproctitle-1.3.3-cp310-cp310-win_amd64.whl", hash = "sha256:200620c3b15388d7f3f97e0ae26599c0c378fdf07ae9ac5a13616e933cbd2086"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:334f7ed39895d692f753a443102dd5fed180c571eb6a48b2a5b7f5b3564908c8"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:950f6476d56ff7817a8fed4ab207727fc5260af83481b2a4b125f32844df513a"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:195c961f54a09eb2acabbfc90c413955cf16c6e2f8caa2adbf2237d1019c7dd8"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f05e66746bf9fe6a3397ec246fe481096664a9c97eb3fea6004735a4daf867fd"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b5901a31012a40ec913265b64e48c2a4059278d9f4e6be628441482dd13fb8b5"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64286f8a995f2cd934082b398fc63fca7d5ffe31f0e27e75b3ca6b4efda4e353"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:184239903bbc6b813b1a8fc86394dc6ca7d20e2ebe6f69f716bec301e4b0199d"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:664698ae0013f986118064b6676d7dcd28fefd0d7d5a5ae9497cbc10cba48fa5"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:e5119a211c2e98ff18b9908ba62a3bd0e3fabb02a29277a7232a6fb4b2560aa0"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:417de6b2e214e837827067048f61841f5d7fc27926f2e43954567094051aff18"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-win32.whl", hash = "sha256:6a143b31d758296dc2f440175f6c8e0b5301ced3b0f477b84ca43cdcf7f2f476"},
|
||||
{file = "setproctitle-1.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:a680d62c399fa4b44899094027ec9a1bdaf6f31c650e44183b50d4c4d0ccc085"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:d4460795a8a7a391e3567b902ec5bdf6c60a47d791c3b1d27080fc203d11c9dc"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:bdfd7254745bb737ca1384dee57e6523651892f0ea2a7344490e9caefcc35e64"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:477d3da48e216d7fc04bddab67b0dcde633e19f484a146fd2a34bb0e9dbb4a1e"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ab2900d111e93aff5df9fddc64cf51ca4ef2c9f98702ce26524f1acc5a786ae7"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:088b9efc62d5aa5d6edf6cba1cf0c81f4488b5ce1c0342a8b67ae39d64001120"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a6d50252377db62d6a0bb82cc898089916457f2db2041e1d03ce7fadd4a07381"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:87e668f9561fd3a457ba189edfc9e37709261287b52293c115ae3487a24b92f6"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:287490eb90e7a0ddd22e74c89a92cc922389daa95babc833c08cf80c84c4df0a"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:4fe1c49486109f72d502f8be569972e27f385fe632bd8895f4730df3c87d5ac8"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4a6ba2494a6449b1f477bd3e67935c2b7b0274f2f6dcd0f7c6aceae10c6c6ba3"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-win32.whl", hash = "sha256:2df2b67e4b1d7498632e18c56722851ba4db5d6a0c91aaf0fd395111e51cdcf4"},
|
||||
{file = "setproctitle-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:f38d48abc121263f3b62943f84cbaede05749047e428409c2c199664feb6abc7"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:816330675e3504ae4d9a2185c46b573105d2310c20b19ea2b4596a9460a4f674"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:68f960bc22d8d8e4ac886d1e2e21ccbd283adcf3c43136161c1ba0fa509088e0"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:00e6e7adff74796ef12753ff399491b8827f84f6c77659d71bd0b35870a17d8f"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:53bc0d2358507596c22b02db079618451f3bd720755d88e3cccd840bafb4c41c"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad6d20f9541f5f6ac63df553b6d7a04f313947f550eab6a61aa758b45f0d5657"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:c1c84beab776b0becaa368254801e57692ed749d935469ac10e2b9b825dbdd8e"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:507e8dc2891021350eaea40a44ddd887c9f006e6b599af8d64a505c0f718f170"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:b1067647ac7aba0b44b591936118a22847bda3c507b0a42d74272256a7a798e9"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:2e71f6365744bf53714e8bd2522b3c9c1d83f52ffa6324bd7cbb4da707312cd8"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-win32.whl", hash = "sha256:7f1d36a1e15a46e8ede4e953abb104fdbc0845a266ec0e99cc0492a4364f8c44"},
|
||||
{file = "setproctitle-1.3.3-cp37-cp37m-win_amd64.whl", hash = "sha256:c9a402881ec269d0cc9c354b149fc29f9ec1a1939a777f1c858cdb09c7a261df"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:ff814dea1e5c492a4980e3e7d094286077054e7ea116cbeda138819db194b2cd"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:accb66d7b3ccb00d5cd11d8c6e07055a4568a24c95cf86109894dcc0c134cc89"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:554eae5a5b28f02705b83a230e9d163d645c9a08914c0ad921df363a07cf39b1"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a911b26264dbe9e8066c7531c0591cfab27b464459c74385b276fe487ca91c12"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2982efe7640c4835f7355fdb4da313ad37fb3b40f5c69069912f8048f77b28c8"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:df3f4274b80709d8bcab2f9a862973d453b308b97a0b423a501bcd93582852e3"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:af2c67ae4c795d1674a8d3ac1988676fa306bcfa1e23fddb5e0bd5f5635309ca"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:af4061f67fd7ec01624c5e3c21f6b7af2ef0e6bab7fbb43f209e6506c9ce0092"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:37a62cbe16d4c6294e84670b59cf7adcc73faafe6af07f8cb9adaf1f0e775b19"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:a83ca086fbb017f0d87f240a8f9bbcf0809f3b754ee01cec928fff926542c450"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-win32.whl", hash = "sha256:059f4ce86f8cc92e5860abfc43a1dceb21137b26a02373618d88f6b4b86ba9b2"},
|
||||
{file = "setproctitle-1.3.3-cp38-cp38-win_amd64.whl", hash = "sha256:ab92e51cd4a218208efee4c6d37db7368fdf182f6e7ff148fb295ecddf264287"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:c7951820b77abe03d88b114b998867c0f99da03859e5ab2623d94690848d3e45"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:5bc94cf128676e8fac6503b37763adb378e2b6be1249d207630f83fc325d9b11"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f5d9027eeda64d353cf21a3ceb74bb1760bd534526c9214e19f052424b37e42"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2e4a8104db15d3462e29d9946f26bed817a5b1d7a47eabca2d9dc2b995991503"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c32c41ace41f344d317399efff4cffb133e709cec2ef09c99e7a13e9f3b9483c"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cbf16381c7bf7f963b58fb4daaa65684e10966ee14d26f5cc90f07049bfd8c1e"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:e18b7bd0898398cc97ce2dfc83bb192a13a087ef6b2d5a8a36460311cb09e775"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:69d565d20efe527bd8a9b92e7f299ae5e73b6c0470f3719bd66f3cd821e0d5bd"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:ddedd300cd690a3b06e7eac90ed4452348b1348635777ce23d460d913b5b63c3"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:415bfcfd01d1fbf5cbd75004599ef167a533395955305f42220a585f64036081"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-win32.whl", hash = "sha256:21112fcd2195d48f25760f0eafa7a76510871bbb3b750219310cf88b04456ae3"},
|
||||
{file = "setproctitle-1.3.3-cp39-cp39-win_amd64.whl", hash = "sha256:5a740f05d0968a5a17da3d676ce6afefebeeeb5ce137510901bf6306ba8ee002"},
|
||||
{file = "setproctitle-1.3.3-pp310-pypy310_pp73-macosx_10_9_x86_64.whl", hash = "sha256:6b9e62ddb3db4b5205c0321dd69a406d8af9ee1693529d144e86bd43bcb4b6c0"},
|
||||
{file = "setproctitle-1.3.3-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9e3b99b338598de0bd6b2643bf8c343cf5ff70db3627af3ca427a5e1a1a90dd9"},
|
||||
{file = "setproctitle-1.3.3-pp310-pypy310_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:38ae9a02766dad331deb06855fb7a6ca15daea333b3967e214de12cfae8f0ef5"},
|
||||
{file = "setproctitle-1.3.3-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:200ede6fd11233085ba9b764eb055a2a191fb4ffb950c68675ac53c874c22e20"},
|
||||
{file = "setproctitle-1.3.3-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:0d3a953c50776751e80fe755a380a64cb14d61e8762bd43041ab3f8cc436092f"},
|
||||
{file = "setproctitle-1.3.3-pp37-pypy37_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e5e08e232b78ba3ac6bc0d23ce9e2bee8fad2be391b7e2da834fc9a45129eb87"},
|
||||
{file = "setproctitle-1.3.3-pp37-pypy37_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f1da82c3e11284da4fcbf54957dafbf0655d2389cd3d54e4eaba636faf6d117a"},
|
||||
{file = "setproctitle-1.3.3-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:aeaa71fb9568ebe9b911ddb490c644fbd2006e8c940f21cb9a1e9425bd709574"},
|
||||
{file = "setproctitle-1.3.3-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:59335d000c6250c35989394661eb6287187854e94ac79ea22315469ee4f4c244"},
|
||||
{file = "setproctitle-1.3.3-pp38-pypy38_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c3ba57029c9c50ecaf0c92bb127224cc2ea9fda057b5d99d3f348c9ec2855ad3"},
|
||||
{file = "setproctitle-1.3.3-pp38-pypy38_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d876d355c53d975c2ef9c4f2487c8f83dad6aeaaee1b6571453cb0ee992f55f6"},
|
||||
{file = "setproctitle-1.3.3-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:224602f0939e6fb9d5dd881be1229d485f3257b540f8a900d4271a2c2aa4e5f4"},
|
||||
{file = "setproctitle-1.3.3-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:d7f27e0268af2d7503386e0e6be87fb9b6657afd96f5726b733837121146750d"},
|
||||
{file = "setproctitle-1.3.3-pp39-pypy39_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f5e7266498cd31a4572378c61920af9f6b4676a73c299fce8ba93afd694f8ae7"},
|
||||
{file = "setproctitle-1.3.3-pp39-pypy39_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33c5609ad51cd99d388e55651b19148ea99727516132fb44680e1f28dd0d1de9"},
|
||||
{file = "setproctitle-1.3.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:eae8988e78192fd1a3245a6f4f382390b61bce6cfcc93f3809726e4c885fa68d"},
|
||||
{file = "setproctitle-1.3.3.tar.gz", hash = "sha256:c913e151e7ea01567837ff037a23ca8740192880198b7fbb90b16d181607caae"},
|
||||
]
|
||||
|
||||
[package.extras]
|
||||
test = ["pytest"]
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "71.0.3"
|
||||
@@ -4823,6 +5067,17 @@ core = ["importlib-metadata (>=6)", "importlib-resources (>=5.10.2)", "jaraco.te
|
||||
doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "pyproject-hooks (!=1.1)", "rst.linker (>=1.9)", "sphinx (<7.4)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (>=1,<2)", "sphinx-reredirects", "sphinxcontrib-towncrier"]
|
||||
test = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "importlib-metadata", "ini2toml[lite] (>=0.14)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "jaraco.test", "mypy (==1.10.0)", "packaging (>=23.2)", "pip (>=19.1)", "pyproject-hooks (!=1.1)", "pytest (>=6,!=8.1.*)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-home (>=0.5)", "pytest-mypy", "pytest-perf", "pytest-ruff (<0.4)", "pytest-ruff (>=0.2.1)", "pytest-ruff (>=0.3.2)", "pytest-subprocess", "pytest-timeout", "pytest-xdist (>=3)", "tomli", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"]
|
||||
|
||||
[[package]]
|
||||
name = "shellingham"
|
||||
version = "1.5.4"
|
||||
description = "Tool to Detect Surrounding Shell"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686"},
|
||||
{file = "shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "six"
|
||||
version = "1.16.0"
|
||||
@@ -4834,6 +5089,17 @@ files = [
|
||||
{file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "smmap"
|
||||
version = "5.0.1"
|
||||
description = "A pure Python implementation of a sliding window memory map manager"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "smmap-5.0.1-py3-none-any.whl", hash = "sha256:e6d8668fa5f93e706934a62d7b4db19c8d9eb8cf2adbb75ef1b675aa332b69da"},
|
||||
{file = "smmap-5.0.1.tar.gz", hash = "sha256:dceeb6c0028fdb6734471eb07c0cd2aae706ccaecab45965ee83f11c8d3b1f62"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sniffio"
|
||||
version = "1.3.1"
|
||||
@@ -5695,6 +5961,23 @@ files = [
|
||||
doc = ["sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
|
||||
test = ["mypy", "pytest", "typing-extensions"]
|
||||
|
||||
[[package]]
|
||||
name = "typer"
|
||||
version = "0.12.3"
|
||||
description = "Typer, build great CLIs. Easy to code. Based on Python type hints."
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "typer-0.12.3-py3-none-any.whl", hash = "sha256:070d7ca53f785acbccba8e7d28b08dcd88f79f1fbda035ade0aecec71ca5c914"},
|
||||
{file = "typer-0.12.3.tar.gz", hash = "sha256:49e73131481d804288ef62598d97a1ceef3058905aa536a1134f90891ba35482"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
click = ">=8.0.0"
|
||||
rich = ">=10.11.0"
|
||||
shellingham = ">=1.3.0"
|
||||
typing-extensions = ">=3.7.4.3"
|
||||
|
||||
[[package]]
|
||||
name = "types-python-dateutil"
|
||||
version = "2.9.0.20240316"
|
||||
@@ -5780,6 +6063,48 @@ dev = ["datasetsforecast (==0.0.8)", "nbdev", "pandas[plot]", "plotly", "plotly-
|
||||
plotting = ["pandas[plot]", "plotly", "plotly-resampler"]
|
||||
polars = ["polars[numpy]"]
|
||||
|
||||
[[package]]
|
||||
name = "wandb"
|
||||
version = "0.17.5"
|
||||
description = "A CLI and library for interacting with the Weights & Biases API."
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "wandb-0.17.5-py3-none-any.whl", hash = "sha256:1c0f60446b51561b67280a060388ffad2a6078fcfdf5024b9998252d237b4639"},
|
||||
{file = "wandb-0.17.5-py3-none-macosx_10_14_x86_64.whl", hash = "sha256:653252c57df550edc70607da827bc68c670932d6775e2f6556909575e17c544b"},
|
||||
{file = "wandb-0.17.5-py3-none-macosx_11_0_arm64.whl", hash = "sha256:233b02d3643142cce8c0ae7986c233fe976b3f7ec0f7aded7478dad0d5a74d43"},
|
||||
{file = "wandb-0.17.5-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ca55edb64f0256a4d4961c3d9dd281a5928037827a21315a8ca67e92ccc60d06"},
|
||||
{file = "wandb-0.17.5-py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:10e4b954ce7ff8615ee64b2dd5a04e51e8c64568d47ec39f5995dbbc9df172db"},
|
||||
{file = "wandb-0.17.5-py3-none-win32.whl", hash = "sha256:04013a6974dd5ff8d69cff79efdbad625db9873e3049bffe85cf39d81f5207cb"},
|
||||
{file = "wandb-0.17.5-py3-none-win_amd64.whl", hash = "sha256:c90e80df09c47e3e0432b2e4e90a4eff34f15e891467ec2f3c284834a33cd6c4"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
click = ">=7.1,<8.0.0 || >8.0.0"
|
||||
docker-pycreds = ">=0.4.0"
|
||||
gitpython = ">=1.0.0,<3.1.29 || >3.1.29"
|
||||
platformdirs = "*"
|
||||
protobuf = {version = ">=3.19.0,<4.21.0 || >4.21.0,<6", markers = "python_version > \"3.9\" or sys_platform != \"linux\""}
|
||||
psutil = ">=5.0.0"
|
||||
pyyaml = "*"
|
||||
requests = ">=2.0.0,<3"
|
||||
sentry-sdk = ">=1.0.0"
|
||||
setproctitle = "*"
|
||||
setuptools = "*"
|
||||
|
||||
[package.extras]
|
||||
aws = ["boto3"]
|
||||
azure = ["azure-identity", "azure-storage-blob"]
|
||||
gcp = ["google-cloud-storage"]
|
||||
importers = ["filelock", "mlflow", "polars", "rich", "tenacity"]
|
||||
kubeflow = ["google-cloud-storage", "kubernetes", "minio", "sh"]
|
||||
launch = ["awscli", "azure-containerregistry", "azure-identity", "azure-storage-blob", "boto3", "botocore", "chardet", "google-auth", "google-cloud-aiplatform", "google-cloud-artifact-registry", "google-cloud-compute", "google-cloud-storage", "iso8601", "kubernetes", "kubernetes-asyncio", "nbconvert", "nbformat", "optuna", "pydantic", "pyyaml (>=6.0.0)", "tomli", "typing-extensions"]
|
||||
media = ["bokeh", "moviepy", "numpy", "pillow", "plotly (>=5.18.0)", "rdkit-pypi", "soundfile"]
|
||||
models = ["cloudpickle"]
|
||||
perf = ["orjson"]
|
||||
sweeps = ["sweeps (>=0.2.0)"]
|
||||
workspaces = ["wandb-workspaces"]
|
||||
|
||||
[[package]]
|
||||
name = "wcwidth"
|
||||
version = "0.2.13"
|
||||
@@ -5972,4 +6297,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = ">=3.10,<3.11"
|
||||
content-hash = "734b625d8c483c4cdced33cc30e90a5199fa1b27724677de68b0429013365853"
|
||||
content-hash = "be8cfad050d901ecd07345980bc91ad5768b0ca000c6c0d741888a6920e5e29e"
|
||||
|
||||
@@ -39,6 +39,11 @@ jax = {version = ">=0.4.26", extras = ["cuda12"]}
|
||||
jaxlib = ">=0.4.26"
|
||||
huggingface_hub = {version = ">=0.23.0", extras = ["cli"]}
|
||||
scikit-learn = ">=1.2.2"
|
||||
typer = ">=0.12.3"
|
||||
wandb = ">=0.17.5"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = ">=8.3.2"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright 2024 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.
|
||||
|
||||
"""adapter init file."""
|
||||
|
||||
from .dora_layers import DoraAttentionProjection, DoraCombinedQKVProjection, DoraLinear
|
||||
from .lora_layers import LoraAttentionProjection, LoraCombinedQKVProjection, LoraLinear
|
||||
@@ -0,0 +1,202 @@
|
||||
# Copyright 2024 The Google Research Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
from jax import numpy as jnp
|
||||
from praxis import base_layer
|
||||
from praxis.layers import attentions, linears
|
||||
|
||||
WeightInit = base_layer.WeightInit
|
||||
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
|
||||
):
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
def _dorafy_var(self, w):
|
||||
lora_a = super().__getattr__("lora_a")
|
||||
lora_b = super().__getattr__("lora_b")
|
||||
dora_m = super().__getattr__("dora_m")
|
||||
|
||||
lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
|
||||
lora_delta = jnp.reshape(lora_delta, w.shape)
|
||||
|
||||
w_prime = w + lora_delta
|
||||
|
||||
column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
|
||||
norm_adapted = w_prime / column_norm
|
||||
w_prime = dora_m * norm_adapted
|
||||
return w_prime
|
||||
|
||||
def __getattr__(self, k):
|
||||
var = super().__getattr__(k)
|
||||
if not self._dora_initialized():
|
||||
return var
|
||||
|
||||
if k == "w":
|
||||
return self._dorafy_var(var)
|
||||
|
||||
return var
|
||||
|
||||
def __getitem__(self, k):
|
||||
var = super().__getattr__(k)
|
||||
if not self._dora_initialized():
|
||||
return var
|
||||
|
||||
if k == "w":
|
||||
return self._dorafy_var(var)
|
||||
|
||||
return var
|
||||
|
||||
|
||||
class DoraThetaDescriptor:
|
||||
"""Dot syntax accession descriptor."""
|
||||
|
||||
def __get__(self, obj, objtype=None):
|
||||
return DoraTheta(obj)
|
||||
|
||||
|
||||
class DoraLinear(linears.Linear):
|
||||
rank: int = 0
|
||||
lora_init: WeightInit | None = None
|
||||
theta = DoraThetaDescriptor()
|
||||
|
||||
def setup(self) -> None:
|
||||
lora_init = self.lora_init if self.lora_init else self.weight_init
|
||||
|
||||
super().setup()
|
||||
self.create_variable(
|
||||
"lora_a",
|
||||
WeightHParams(
|
||||
shape=[self.input_dims, self.rank],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"lora_b",
|
||||
WeightHParams(
|
||||
shape=[self.output_dims, self.rank],
|
||||
init=WeightInit.Constant(scale=0.0),
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"dora_m",
|
||||
WeightHParams(
|
||||
shape=[1, self.output_dims],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class DoraAttentionProjection(attentions.AttentionProjection):
|
||||
rank: int = 0
|
||||
lora_init: WeightInit | None = None
|
||||
theta = DoraThetaDescriptor()
|
||||
|
||||
def setup(self) -> None:
|
||||
super().setup()
|
||||
w_weight_params = self._weight_hparams["w"]
|
||||
lora_init = self.lora_init if self.lora_init else w_weight_params.init
|
||||
|
||||
self.create_variable(
|
||||
"lora_a",
|
||||
WeightHParams(
|
||||
shape=[self.input_dim, self.rank],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[
|
||||
None,
|
||||
None,
|
||||
],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"lora_b",
|
||||
WeightHParams(
|
||||
shape=[self.dim_per_head * self.num_heads, self.rank],
|
||||
init=WeightInit.Constant(scale=0.0),
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[
|
||||
None,
|
||||
None,
|
||||
],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"dora_m",
|
||||
WeightHParams(
|
||||
shape=[1, self.num_heads, self.dim_per_head],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None, None],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class DoraCombinedQKVProjection(attentions.CombinedQKVProjectionLayer):
|
||||
rank: int = 0
|
||||
lora_init: WeightInit | None = None
|
||||
theta = DoraThetaDescriptor()
|
||||
|
||||
def setup(self) -> None:
|
||||
super().setup()
|
||||
w_weight_params = self._weight_hparams["w"]
|
||||
lora_init = self.lora_init if self.lora_init else w_weight_params.init
|
||||
|
||||
self.create_variable(
|
||||
"lora_a",
|
||||
WeightHParams(
|
||||
shape=[3, self.input_dim, self.rank],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None, None],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"lora_b",
|
||||
WeightHParams(
|
||||
shape=[3, self.dim_per_head * self.num_heads, self.rank],
|
||||
init=WeightInit.Constant(scale=0.0),
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None, None],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"dora_m",
|
||||
WeightHParams(
|
||||
shape=[3, 1, self.num_heads, self.dim_per_head],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None, None, None],
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,166 @@
|
||||
# Copyright 2024 The Google Research Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
from jax import numpy as jnp
|
||||
from praxis import base_layer
|
||||
from praxis.layers import attentions, linears
|
||||
|
||||
WeightInit = base_layer.WeightInit
|
||||
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
|
||||
):
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
def _lorafy_var(self, w):
|
||||
lora_a = super().__getattr__("lora_a")
|
||||
lora_b = super().__getattr__("lora_b")
|
||||
lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
|
||||
lora_delta = jnp.reshape(lora_delta, w.shape)
|
||||
w_prime = w + lora_delta
|
||||
return w_prime
|
||||
|
||||
def __getattr__(self, k):
|
||||
var = super().__getattr__(k)
|
||||
if not self._lora_initialized():
|
||||
return var
|
||||
|
||||
if k == "w":
|
||||
return self._lorafy_var(var)
|
||||
|
||||
return var
|
||||
|
||||
def __getitem__(self, k):
|
||||
var = super().__getattr__(k)
|
||||
if not self._lora_initialized():
|
||||
return var
|
||||
|
||||
if k == "w":
|
||||
return self._lorafy_var(var)
|
||||
|
||||
return var
|
||||
|
||||
|
||||
class LoraThetaDescriptor:
|
||||
"""Dot syntax accession descriptor."""
|
||||
|
||||
def __get__(self, obj, objtype=None):
|
||||
return LoraTheta(obj)
|
||||
|
||||
|
||||
class LoraLinear(linears.Linear):
|
||||
rank: int = 0
|
||||
lora_init: WeightInit | None = None
|
||||
theta = LoraThetaDescriptor()
|
||||
|
||||
def setup(self) -> None:
|
||||
lora_init = self.lora_init if self.lora_init else self.weight_init
|
||||
|
||||
super().setup()
|
||||
self.create_variable(
|
||||
"lora_a",
|
||||
WeightHParams(
|
||||
shape=[self.input_dims, self.rank],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"lora_b",
|
||||
WeightHParams(
|
||||
shape=[self.output_dims, self.rank],
|
||||
init=WeightInit.Constant(scale=0.0),
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class LoraAttentionProjection(attentions.AttentionProjection):
|
||||
rank: int = 0
|
||||
lora_init: WeightInit | None = None
|
||||
theta = LoraThetaDescriptor()
|
||||
|
||||
def setup(self) -> None:
|
||||
super().setup()
|
||||
w_weight_params = self._weight_hparams["w"]
|
||||
lora_init = self.lora_init if self.lora_init else w_weight_params.init
|
||||
|
||||
self.create_variable(
|
||||
"lora_a",
|
||||
WeightHParams(
|
||||
shape=[self.input_dim, self.rank],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[
|
||||
None,
|
||||
None,
|
||||
],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"lora_b",
|
||||
WeightHParams(
|
||||
shape=[self.dim_per_head * self.num_heads, self.rank],
|
||||
init=WeightInit.Constant(scale=0.0),
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[
|
||||
None,
|
||||
None,
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class LoraCombinedQKVProjection(attentions.CombinedQKVProjectionLayer):
|
||||
rank: int = 0
|
||||
lora_init: WeightInit | None = None
|
||||
theta = LoraThetaDescriptor()
|
||||
|
||||
def setup(self) -> None:
|
||||
super().setup()
|
||||
w_weight_params = self._weight_hparams["w"]
|
||||
lora_init = self.lora_init if self.lora_init else w_weight_params.init
|
||||
|
||||
self.create_variable(
|
||||
"lora_a",
|
||||
WeightHParams(
|
||||
shape=[3, self.input_dim, self.rank],
|
||||
init=lora_init,
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None, None],
|
||||
),
|
||||
)
|
||||
self.create_variable(
|
||||
"lora_b",
|
||||
WeightHParams(
|
||||
shape=[3, self.dim_per_head * self.num_heads, self.rank],
|
||||
init=WeightInit.Constant(scale=0.0),
|
||||
mesh_shape=self.mesh_shape,
|
||||
tensor_split_dims_mapping=[None, None, None],
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,487 @@
|
||||
# Copyright 2024 The Google Research Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
"""
|
||||
This file provides functionality for loading and merging adapter weights
|
||||
in timesfm model, specifically for LoRA and DoRA.
|
||||
LoRA: https://arxiv.org/abs/2106.09685
|
||||
DoRA: https://arxiv.org/abs/2402.09353v4
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
from paxml import checkpoints, tasks_lib
|
||||
from paxml.train_states import TrainState
|
||||
from praxis import pax_fiddle
|
||||
|
||||
from adapter.dora_layers import (
|
||||
DoraAttentionProjection,
|
||||
DoraCombinedQKVProjection,
|
||||
DoraLinear,
|
||||
)
|
||||
from adapter.lora_layers import (
|
||||
LoraAttentionProjection,
|
||||
LoraCombinedQKVProjection,
|
||||
LoraLinear,
|
||||
)
|
||||
from timesfm import TimesFm
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
params (dict): The full model parameters.
|
||||
lora_target_modules (str): Target modules for LoRA/DoRA adaptation.
|
||||
num_layers (int): Number of transformer layers.
|
||||
use_dora (bool, optional): Whether DoRA was used or not. Defaults to False.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the extracted adapter parameters.
|
||||
"""
|
||||
adapter_params = {}
|
||||
for i in range(num_layers):
|
||||
layer_key = f"x_layers_{i}"
|
||||
adapter_params[layer_key] = {}
|
||||
|
||||
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"]
|
||||
|
||||
lora_a = linear["lora_a"]
|
||||
lora_b = linear["lora_b"]
|
||||
|
||||
adapter_params[layer_key][ff_layer_key] = {
|
||||
"lora_a": lora_a,
|
||||
"lora_b": lora_b,
|
||||
}
|
||||
|
||||
if use_dora:
|
||||
adapter_params[layer_key][ff_layer_key]["dora_m"] = linear["dora_m"]
|
||||
|
||||
if lora_target_modules in ["all", "attention"]:
|
||||
attention = params["params"]["core_layer"]["stacked_transformer_layer"][
|
||||
layer_key
|
||||
]["self_attention"]
|
||||
|
||||
for component in ["key", "query", "value", "post"]:
|
||||
lora_a = attention[component]["lora_a"]
|
||||
lora_b = attention[component]["lora_b"]
|
||||
|
||||
adapter_params[layer_key][component] = {
|
||||
"lora_a": lora_a,
|
||||
"lora_b": lora_b,
|
||||
}
|
||||
|
||||
if use_dora:
|
||||
adapter_params[layer_key][component]["dora_m"] = attention[
|
||||
component
|
||||
]["dora_m"]
|
||||
return adapter_params
|
||||
|
||||
|
||||
def load_adapter_checkpoint(
|
||||
model: TimesFm,
|
||||
adapter_checkpoint_path: str,
|
||||
lora_rank: int,
|
||||
lora_target_modules: str,
|
||||
use_dora: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Loads an adapter checkpoint and merges it with the original model weights.
|
||||
|
||||
Args:
|
||||
model (TimesFm): The model to update.
|
||||
adapter_checkpoint_path (str): Path to the adapter checkpoint.
|
||||
lora_rank (int): Rank of the LoRA adaptation.
|
||||
lora_target_modules (str): Target modules for adaptation.
|
||||
use_dora (bool): Whether DoRA was used or not.
|
||||
|
||||
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}.")
|
||||
start_time = time.time()
|
||||
original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl = (
|
||||
load_adapter_layer(
|
||||
mdl_vars=model._train_state.mdl_vars,
|
||||
model=model._model,
|
||||
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
|
||||
)
|
||||
|
||||
adapter_weight_hparams = _get_adapter_weight_params(
|
||||
var_weight_hparams=var_weight_hparams,
|
||||
lora_target_modules=lora_target_modules,
|
||||
num_layers=model._model.stacked_transformer_params_tpl.num_layers,
|
||||
use_dora=use_dora,
|
||||
)
|
||||
|
||||
adapter_state_partition_specs = tasks_lib.create_state_partition_specs(
|
||||
adapter_weight_hparams,
|
||||
mesh_shape=model.mesh_shape,
|
||||
mesh_axis_names=model.mesh_name,
|
||||
discard_opt_states=True,
|
||||
learners=None,
|
||||
)
|
||||
adapter_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
|
||||
adapter_weight_hparams,
|
||||
discard_opt_states=True,
|
||||
learners=None,
|
||||
)
|
||||
adapter_train_state = checkpoints.restore_checkpoint(
|
||||
state_global_shapes=adapter_state_local_shapes,
|
||||
checkpoint_dir=adapter_checkpoint_path,
|
||||
checkpoint_type=checkpoints.CheckpointType.FLAX,
|
||||
state_specs=adapter_state_partition_specs,
|
||||
step=None,
|
||||
)
|
||||
|
||||
# add adapter weights to the original weights
|
||||
_merge_adapter_weights(
|
||||
model=model,
|
||||
adapter_train_state=adapter_train_state,
|
||||
lora_target_modules=lora_target_modules,
|
||||
num_layers=model._model.stacked_transformer_params_tpl.num_layers,
|
||||
use_dora=use_dora,
|
||||
)
|
||||
|
||||
# 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
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
|
||||
original_combined_qkv_tpl
|
||||
)
|
||||
model._logging(
|
||||
f"Restored adapter checkpoint in {time.time() - start_time:.2f} seconds."
|
||||
)
|
||||
|
||||
# jit compile the model
|
||||
model.jit_decode()
|
||||
|
||||
|
||||
def _merge_adapter_weights(
|
||||
model: TimesFm,
|
||||
adapter_train_state: TrainState,
|
||||
lora_target_modules: str,
|
||||
num_layers: int,
|
||||
use_dora: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Merges adapter weights with the original model weights.
|
||||
|
||||
Args:
|
||||
model (TimesFm): The model to update.
|
||||
adapter_train_state (TrainState): The adapter's train state.
|
||||
lora_target_modules (str): Target modules for adaptation.
|
||||
num_layers (int): Number of transformer layers.
|
||||
use_dora (bool): Whether DoRA was used or not.
|
||||
"""
|
||||
for i in range(num_layers):
|
||||
layer_key = f"x_layers_{i}"
|
||||
|
||||
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"]
|
||||
|
||||
params = adapter_train_state.mdl_vars[layer_key][ff_layer_key]
|
||||
lora_a = params["lora_a"]
|
||||
lora_b = params["lora_b"]
|
||||
|
||||
w = linear["w"]
|
||||
|
||||
lora_delta = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
|
||||
lora_delta = jnp.reshape(lora_delta, w.shape)
|
||||
w_prime = w + lora_delta
|
||||
|
||||
if use_dora:
|
||||
dora_m = params["dora_m"]
|
||||
column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
|
||||
norm_adapted = w_prime / column_norm
|
||||
w_prime = dora_m * norm_adapted
|
||||
linear["w"] = w_prime
|
||||
del linear["dora_m"]
|
||||
|
||||
else:
|
||||
linear["w"] = w_prime
|
||||
|
||||
del linear["lora_a"]
|
||||
del linear["lora_b"]
|
||||
|
||||
if lora_target_modules in ["all", "attention"]:
|
||||
attention = model._train_state.mdl_vars["params"][
|
||||
"stacked_transformer_layer"
|
||||
][layer_key]["self_attention"]
|
||||
|
||||
for component in ["key", "query", "value", "post"]:
|
||||
params = adapter_train_state.mdl_vars[layer_key][component]
|
||||
lora_a = params["lora_a"]
|
||||
lora_b = params["lora_b"]
|
||||
|
||||
w = attention[component]["w"]
|
||||
|
||||
lora_delta = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
|
||||
lora_delta = jnp.reshape(lora_delta, w.shape)
|
||||
w_prime = w + lora_delta
|
||||
|
||||
if use_dora:
|
||||
dora_m = params["dora_m"]
|
||||
column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
|
||||
norm_adapted = w_prime / column_norm
|
||||
w_prime = dora_m * norm_adapted
|
||||
attention[component]["w"] = w_prime
|
||||
del attention[component]["dora_m"]
|
||||
|
||||
else:
|
||||
attention[component]["w"] = w_prime
|
||||
|
||||
del attention[component]["lora_a"]
|
||||
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:
|
||||
"""
|
||||
Extracts adapter weight parameters from the given variable weight hyperparameters.
|
||||
|
||||
Args:
|
||||
var_weight_hparams (dict): Variable weight hyperparameters.
|
||||
lora_target_modules (str): Target modules for adaptation.
|
||||
num_layers (int): Number of transformer layers.
|
||||
use_dora (bool): Whether DoRA was used or not.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the extracted adapter weight parameters.
|
||||
"""
|
||||
adapter_params = {}
|
||||
for i in range(num_layers):
|
||||
layer = f"x_layers_{i}"
|
||||
adapter_params[layer] = {}
|
||||
|
||||
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"]
|
||||
adapter_params[layer][ff_layer_key] = {
|
||||
"lora_a": adapter_weight_params["lora_a"],
|
||||
"lora_b": adapter_weight_params["lora_b"],
|
||||
}
|
||||
|
||||
if use_dora:
|
||||
adapter_params[layer][ff_layer_key]["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]
|
||||
adapter_params[layer][component] = {
|
||||
"lora_a": adapter_weight_params["lora_a"],
|
||||
"lora_b": adapter_weight_params["lora_b"],
|
||||
}
|
||||
|
||||
if use_dora:
|
||||
adapter_params[layer][component]["dora_m"] = adapter_weight_params[
|
||||
"dora_m"
|
||||
]
|
||||
|
||||
return adapter_params
|
||||
|
||||
|
||||
def load_adapter_layer(
|
||||
mdl_vars: dict,
|
||||
model: pax_fiddle.Config,
|
||||
lora_rank: int,
|
||||
lora_target_modules: str,
|
||||
use_dora: bool = False,
|
||||
) -> tuple[pax_fiddle.Config, pax_fiddle.Config]:
|
||||
"""
|
||||
Updates target modules with adapter layers.
|
||||
|
||||
Args:
|
||||
mdl_vars (dict): Model variables.
|
||||
model (pax_fiddle.Config): Model configuration.
|
||||
lora_rank (int): Rank of the LoRA adaptation.
|
||||
lora_target_modules (str): Target modules for adaptation.
|
||||
use_dora (bool, optional): Whether DoRA was used or not.
|
||||
|
||||
Returns:
|
||||
tuple[pax_fiddle.Config, pax_fiddle.Config]: Updated model configurations.
|
||||
"""
|
||||
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(
|
||||
DoraLinear,
|
||||
rank=lora_rank,
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
if lora_target_modules in ["all", "attention"]:
|
||||
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)
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
adapter_combined_qkv_tpl = (
|
||||
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
|
||||
)
|
||||
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
|
||||
adapter_combined_qkv_tpl
|
||||
)
|
||||
|
||||
# initialize and add adapter weights
|
||||
_initialize_adapter_params(
|
||||
mdl_vars=mdl_vars,
|
||||
num_layers=model.stacked_transformer_params_tpl.num_layers,
|
||||
lora_rank=lora_rank,
|
||||
lora_target_modules=lora_target_modules,
|
||||
use_dora=use_dora,
|
||||
)
|
||||
|
||||
return original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl
|
||||
|
||||
|
||||
def _initialize_adapter_params(
|
||||
mdl_vars: dict,
|
||||
num_layers,
|
||||
lora_rank: int,
|
||||
lora_target_modules: str,
|
||||
use_dora: bool = False,
|
||||
seed: int = 1234,
|
||||
) -> dict:
|
||||
"""
|
||||
Initializes and adds adapter parameters to target modules.
|
||||
|
||||
Args:
|
||||
mdl_vars (dict): Model variables.
|
||||
num_layers (int): Number of transformer layers.
|
||||
lora_rank (int): Rank of the LoRA adaptation.
|
||||
lora_target_modules (str): Target modules for adaptation.
|
||||
use_dora (bool, optional): Whether DoRA was used or not.
|
||||
seed (int, optional): Random seed for initialization. Defaults to 1234.
|
||||
|
||||
Returns:
|
||||
dict: Updated model variables with initialized adapter parameters.
|
||||
"""
|
||||
for i in range(num_layers):
|
||||
layer_key = f"x_layers_{i}"
|
||||
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"]
|
||||
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_b = jnp.zeros((output_dim, lora_rank))
|
||||
|
||||
linear["lora_a"] = lora_a
|
||||
linear["lora_b"] = lora_b
|
||||
|
||||
if use_dora:
|
||||
norm = jnp.linalg.norm(original_w, ord=2, axis=0, keepdims=True)
|
||||
linear["dora_m"] = norm
|
||||
|
||||
if lora_target_modules in ["all", "attention"]:
|
||||
attention = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
|
||||
"self_attention"
|
||||
]
|
||||
|
||||
for component in ["key", "query", "value", "post"]:
|
||||
original_w = attention[component]["w"]
|
||||
w_dim = original_w.shape[0]
|
||||
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_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)
|
||||
attention[component]["dora_m"] = norm
|
||||
return mdl_vars
|
||||
@@ -0,0 +1,91 @@
|
||||
# Copyright 2024 The Google Research Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import timesfm
|
||||
|
||||
|
||||
def create_sample_dataframe(
|
||||
start_date: datetime, end_date: datetime, freq: str = "D"
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Create a sample DataFrame with time series data.
|
||||
|
||||
Args:
|
||||
start_date (datetime): Start date of the time series.
|
||||
end_date (datetime): End date of the time series.
|
||||
freq (str): Frequency of the time series (default: "D" for daily).
|
||||
|
||||
Returns:
|
||||
pd.DataFrame: DataFrame with columns 'unique_id', 'ds', and 'ts'.
|
||||
"""
|
||||
date_range = pd.date_range(start=start_date, end=end_date, freq=freq)
|
||||
ts_data = np.random.randn(len(date_range))
|
||||
df = pd.DataFrame({"unique_id": "ts-1", "ds": date_range, "ts": ts_data})
|
||||
return df
|
||||
|
||||
|
||||
@pytest.mark.parametrize("context_length", [128, 256, 512])
|
||||
@pytest.mark.parametrize("prediction_length", [96, 128, 256])
|
||||
@pytest.mark.parametrize("freq", ["D", "H", "W"])
|
||||
def test_timesfm_forecast_on_df(
|
||||
context_length: int,
|
||||
prediction_length: int,
|
||||
freq: str,
|
||||
) -> None:
|
||||
model = timesfm.TimesFm(
|
||||
context_len=context_length,
|
||||
horizon_len=prediction_length,
|
||||
input_patch_len=32,
|
||||
output_patch_len=128,
|
||||
num_layers=20,
|
||||
model_dims=1280,
|
||||
backend="cpu",
|
||||
)
|
||||
model.load_from_checkpoint(repo_id="google/timesfm-1.0-200m")
|
||||
|
||||
end_date = datetime.now()
|
||||
start_date = end_date - timedelta(days=context_length)
|
||||
input_df = create_sample_dataframe(start_date, end_date, freq)
|
||||
|
||||
forecast_df = model.forecast_on_df(
|
||||
inputs=input_df,
|
||||
freq=freq,
|
||||
value_name="ts",
|
||||
num_jobs=-1,
|
||||
)
|
||||
|
||||
assert (
|
||||
len(forecast_df) == prediction_length
|
||||
), f"Expected forecast length of {prediction_length}, but got {len(forecast_df)}"
|
||||
assert (
|
||||
"timesfm" in forecast_df.columns
|
||||
), "Forecast DataFrame should contain 'timesfm' column"
|
||||
|
||||
last_input_date = input_df["ds"].max()
|
||||
first_forecast_date = forecast_df["ds"].min()
|
||||
expected_first_forecast_date = last_input_date + pd.Timedelta(1, unit=freq)
|
||||
assert (
|
||||
first_forecast_date == expected_first_forecast_date
|
||||
), f"Forecast should start from {expected_first_forecast_date}, but starts from {first_forecast_date}"
|
||||
|
||||
print(
|
||||
f"Successful forecast with context_length={context_length}, prediction_length={prediction_length}, freq={freq}"
|
||||
)
|
||||
Reference in New Issue
Block a user