diff --git a/.gitignore b/.gitignore index cc38b4f..c39ae71 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,8 @@ .venv/ dist/ -**__pycache__/** */ \ No newline at end of file +__pycache__/ +checkpoints/ +wandb/ +datasets/ +results/ +timesfm_jax.egg-info/ diff --git a/environment.yml b/environment.yml index 7a1d6fe..a28abf6 100644 --- a/environment.yml +++ b/environment.yml @@ -16,3 +16,6 @@ dependencies: - jax[cuda12]==0.4.26 - einshape - scikit-learn + - typer + - wandb + - pytest diff --git a/environment_cpu.yml b/environment_cpu.yml index d808176..c539642 100644 --- a/environment_cpu.yml +++ b/environment_cpu.yml @@ -16,3 +16,6 @@ dependencies: - jax[cpu]==0.4.26 - einshape - scikit-learn + - typer + - wandb + - pytest diff --git a/peft/README.md b/peft/README.md new file mode 100644 index 0000000..8730f03 --- /dev/null +++ b/peft/README.md @@ -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: + +image + +- 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. diff --git a/peft/finetune.py b/peft/finetune.py new file mode 100644 index 0000000..84c59f9 --- /dev/null +++ b/peft/finetune.py @@ -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) diff --git a/peft/finetune.sh b/peft/finetune.sh new file mode 100644 index 0000000..3ada1c3 --- /dev/null +++ b/peft/finetune.sh @@ -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 diff --git a/peft/usage.ipynb b/peft/usage.ipynb new file mode 100644 index 0000000..ab5177a --- /dev/null +++ b/peft/usage.ipynb @@ -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 +} diff --git a/poetry.lock b/poetry.lock index 9131807..a33d8c1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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" diff --git a/pyproject.toml b/pyproject.toml index 2a56af8..48d08dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/adapter/__init__.py b/src/adapter/__init__.py new file mode 100644 index 0000000..6870b01 --- /dev/null +++ b/src/adapter/__init__.py @@ -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 diff --git a/src/adapter/dora_layers.py b/src/adapter/dora_layers.py new file mode 100644 index 0000000..9a28911 --- /dev/null +++ b/src/adapter/dora_layers.py @@ -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], + ), + ) diff --git a/src/adapter/lora_layers.py b/src/adapter/lora_layers.py new file mode 100644 index 0000000..15df5a5 --- /dev/null +++ b/src/adapter/lora_layers.py @@ -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], + ), + ) diff --git a/src/adapter/utils.py b/src/adapter/utils.py new file mode 100644 index 0000000..4c3fc5b --- /dev/null +++ b/src/adapter/utils.py @@ -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 diff --git a/tests/test_timesfm.py b/tests/test_timesfm.py new file mode 100644 index 0000000..3277a9a --- /dev/null +++ b/tests/test_timesfm.py @@ -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}" + )