diff --git a/README.md b/README.md index 5549b0c..c5a1c43 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,40 @@ To install the TimesFM as a package, you can run the following command without c `pip install timesfm` -### Local installation +### Installation using conda + +For calling TimesFM, We have two environment files. Inside `timesfm`, for +GPU installation (assuming CUDA 12 has been setup), you can create a conda +environment `tfm_env` from the base folder through: + +``` +conda env create --file=environment.yml +``` + +For a CPU setup please use, + +``` +conda env create --file=environment_cpu.yml +``` +to create the environment instead. + +Follow by + +``` +conda activate tfm_env +pip install -e . +``` +to install the package. + +**Note**: + +1. Running the provided benchmarks would require additional dependencies. +Please use the environment files under `experiments` instead. + +2. The dependency `lingvo` does not support ARM architectures, and the code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned. + + +### Local installation using poetry To from the current repository/local version (like you would have previously done with `pip -e .`), you can run the command @@ -49,7 +82,7 @@ poetry install This will install the environment in the local .venv folder (depends on the configuration) and matches the python command to the poetry environment. If this is not the case, you can use `poetry run python` to use the local environment. -## Notes +### Notes 1. Running the provided benchmarks would require additional dependencies. Please use the environment files under `experiments` instead. diff --git a/poetry.lock b/poetry.lock index 997979b..11dc0f7 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1620,8 +1620,20 @@ files = [ ] [package.dependencies] +jax-cuda12-plugin = {version = "0.4.26", optional = true, markers = "extra == \"cuda12\""} +jaxlib = {version = "0.4.26", optional = true, markers = "extra == \"cuda12\""} ml-dtypes = ">=0.2.0" numpy = ">=1.22" +nvidia-cublas-cu12 = {version = ">=12.1.3.1", optional = true, markers = "extra == \"cuda12\""} +nvidia-cuda-cupti-cu12 = {version = ">=12.1.105", optional = true, markers = "extra == \"cuda12\""} +nvidia-cuda-nvcc-cu12 = {version = ">=12.1.105", optional = true, markers = "extra == \"cuda12\""} +nvidia-cuda-runtime-cu12 = {version = ">=12.1.105", optional = true, markers = "extra == \"cuda12\""} +nvidia-cudnn-cu12 = {version = ">=8.9.2.26,<9.0", optional = true, markers = "extra == \"cuda12\""} +nvidia-cufft-cu12 = {version = ">=11.0.2.54", optional = true, markers = "extra == \"cuda12\""} +nvidia-cusolver-cu12 = {version = ">=11.4.5.107", optional = true, markers = "extra == \"cuda12\""} +nvidia-cusparse-cu12 = {version = ">=12.1.0.106", optional = true, markers = "extra == \"cuda12\""} +nvidia-nccl-cu12 = {version = ">=2.18.1", optional = true, markers = "extra == \"cuda12\""} +nvidia-nvjitlink-cu12 = {version = ">=12.1.105", optional = true, markers = "extra == \"cuda12\""} opt-einsum = "*" scipy = ">=1.9" @@ -1648,6 +1660,32 @@ files = [ {file = "jax_bitempered_loss-0.0.2-py3-none-any.whl", hash = "sha256:480c8b80abcb00883b001fbdd6f3f70b27681bd5e204fba6669818c0f9a326f1"}, ] +[[package]] +name = "jax-cuda12-pjrt" +version = "0.4.26" +description = "JAX XLA PJRT Plugin for NVIDIA GPUs" +optional = false +python-versions = "*" +files = [ + {file = "jax_cuda12_pjrt-0.4.26-py3-none-manylinux2014_x86_64.whl", hash = "sha256:05e4cca6bc0159ea931a560481b573f4cfbaa700e4e0db8f9a851ffc8e60f450"}, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.4.26" +description = "JAX Plugin for NVIDIA GPUs" +optional = false +python-versions = ">=3.9" +files = [ + {file = "jax_cuda12_plugin-0.4.26-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:50c958cea070dc9598d4bc4b514c674b6e08123fc70ffeeb3dac0fbed9270322"}, + {file = "jax_cuda12_plugin-0.4.26-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:9a857dd23758a06630355d616388e1b74fc8083e55f14473555422ead65ba3c4"}, + {file = "jax_cuda12_plugin-0.4.26-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:7bf1c1044f72c6ad80cb4347bf3c959b1b409fba21e669a1fe3ae4b58c19a6ba"}, + {file = "jax_cuda12_plugin-0.4.26-cp39-cp39-manylinux2014_x86_64.whl", hash = "sha256:79d4fd24b817d507cc6c0c14e88f483038335ef024202a29eb3dd6df583741bc"}, +] + +[package.dependencies] +jax-cuda12-pjrt = "0.4.26" + [[package]] name = "jaxlib" version = "0.4.26" @@ -1955,13 +1993,13 @@ jupyter-server = ">=1.1.2" [[package]] name = "jupyter-server" -version = "2.14.1" +version = "2.14.2" description = "The backend—i.e. core services, APIs, and REST endpoints—to Jupyter web applications." optional = false python-versions = ">=3.8" files = [ - {file = "jupyter_server-2.14.1-py3-none-any.whl", hash = "sha256:16f7177c3a4ea8fe37784e2d31271981a812f0b2874af17339031dc3510cc2a5"}, - {file = "jupyter_server-2.14.1.tar.gz", hash = "sha256:12558d158ec7a0653bf96cc272bc7ad79e0127d503b982ed144399346694f726"}, + {file = "jupyter_server-2.14.2-py3-none-any.whl", hash = "sha256:47ff506127c2f7851a17bf4713434208fc490955d0e8632e95014a9a9afbeefd"}, + {file = "jupyter_server-2.14.2.tar.gz", hash = "sha256:66095021aa9638ced276c248b1d81862e4c50f292d575920bbe960de1c56b12b"}, ] [package.dependencies] @@ -3041,6 +3079,141 @@ files = [ {file = "numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010"}, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.5.3.2" +description = "CUBLAS native runtime libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cublas_cu12-12.5.3.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b"}, + {file = "nvidia_cublas_cu12-12.5.3.2-py3-none-win_amd64.whl", hash = "sha256:4960f3dc5f39699acadf76fa6d94b10a2a00f2956c2c442efa299fb22b0748f3"}, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.5.82" +description = "CUDA profiling tools runtime libs." +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_cupti_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953"}, + {file = "nvidia_cuda_cupti_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3"}, +] + +[[package]] +name = "nvidia-cuda-nvcc-cu12" +version = "12.5.82" +description = "CUDA nvcc" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_nvcc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0"}, + {file = "nvidia_cuda_nvcc_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957"}, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.5.82" +description = "NVRTC native runtime libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_nvrtc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd"}, + {file = "nvidia_cuda_nvrtc_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78"}, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.5.82" +description = "CUDA Runtime native Libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_runtime_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115"}, + {file = "nvidia_cuda_runtime_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46"}, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "8.9.7.29" +description = "cuDNN runtime libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cudnn_cu12-8.9.7.29-py3-none-manylinux1_x86_64.whl", hash = "sha256:adf4f59ed7a1341103822ed8df6e144f4d47ea8b10d9bf0ea0047ba738fd7b02"}, + {file = "nvidia_cudnn_cu12-8.9.7.29-py3-none-win_amd64.whl", hash = "sha256:4447321a2bdc8bd965084c1824575eb04f47a03ab62bbeb6cce7e9f74c3657f3"}, +] + +[package.dependencies] +nvidia-cublas-cu12 = "*" +nvidia-cuda-nvrtc-cu12 = "*" + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.2.3.61" +description = "CUFFT native runtime libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cufft_cu12-11.2.3.61-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8"}, + {file = "nvidia_cufft_cu12-11.2.3.61-py3-none-win_amd64.whl", hash = "sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b"}, +] + +[package.dependencies] +nvidia-nvjitlink-cu12 = "*" + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.6.3.83" +description = "CUDA solver native runtime libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cusolver_cu12-11.6.3.83-py3-none-manylinux2014_x86_64.whl", hash = "sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b"}, + {file = "nvidia_cusolver_cu12-11.6.3.83-py3-none-win_amd64.whl", hash = "sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf"}, +] + +[package.dependencies] +nvidia-cublas-cu12 = "*" +nvidia-cusparse-cu12 = "*" +nvidia-nvjitlink-cu12 = "*" + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.1.3" +description = "CUSPARSE native runtime libraries" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_cusparse_cu12-12.5.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752"}, + {file = "nvidia_cusparse_cu12-12.5.1.3-py3-none-win_amd64.whl", hash = "sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1"}, +] + +[package.dependencies] +nvidia-nvjitlink-cu12 = "*" + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.22.3" +description = "NVIDIA Collective Communication Library (NCCL) Runtime" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_nccl_cu12-2.22.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f9f5e03c00269dee2cd1aa57019f9a024478a74ae6e9b32d5341c849fe6f6302"}, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.5.82" +description = "Nvidia JIT LTO Library" +optional = false +python-versions = ">=3" +files = [ + {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212"}, + {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697"}, +] + [[package]] name = "oauthlib" version = "3.2.2" @@ -5782,4 +5955,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.11" -content-hash = "5af00562d4d4673c273707e744e7cb7ab12b7cbb003233b5d434bc9aeafb4538" +content-hash = "b4a9878c96695bc88213e12bac34aaf4c95175e22c83b65ffaf522096ac32afb" diff --git a/pyproject.toml b/pyproject.toml index e3c63f0..d8b0f3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ numpy = "1.26.4" pandas = "2.1.4" paxml = "1.4.0" utilsforecast = "0.1.10" -jax = "0.4.26" +jax = {version = "0.4.26", extras = ["cuda12"]} jaxlib = "0.4.26" huggingface_hub = {version = "0.23.0", extras = ["cli"]} scikit-learn = "1.5.1" diff --git a/src/timesfm/data_loader.py b/src/timesfm/data_loader.py index eeace7b..d81b130 100644 --- a/src/timesfm/data_loader.py +++ b/src/timesfm/data_loader.py @@ -11,13 +11,11 @@ # 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. - """TF dataloaders for general timeseries datasets. The expected input format is csv file with a datetime index. """ - from absl import logging import numpy as np import pandas as pd @@ -79,9 +77,8 @@ class TimeSeriesdata(object): self.data_df['ccol'] = np.zeros(self.data_df.shape[0]) cat_cov_cols = ['ccol'] self.data_df.fillna(0, inplace=True) - self.data_df.set_index( - pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True - ) + self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]), + inplace=True) self.num_cov_cols = num_cov_cols self.cat_cov_cols = cat_cov_cols self.ts_cols = ts_cols @@ -94,18 +91,16 @@ class TimeSeriesdata(object): data_df_idx[-1] + pd.Timedelta(1, freq=freq), periods=pred_len + 1, freq=freq, - ) - ) + )) self.time_df = time_features.TimeCovariates( - date_index, holiday=holiday - ).get_covariates() + date_index, holiday=holiday).get_covariates() self.hist_len = hist_len self.pred_len = pred_len self.batch_size = batch_size self.freq = freq self.normalize = normalize self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose() - self.data_mat = self.data_mat[:, 0 : self.test_range[1]] + self.data_mat = self.data_mat[:, 0:self.test_range[1]] self.time_mat = self.time_df.to_numpy().transpose() self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose() self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols) @@ -135,7 +130,7 @@ class TimeSeriesdata(object): def _normalize_data(self): self.scaler = StandardScaler() - train_mat = self.data_mat[:, self.train_range[0] : self.train_range[1]] + train_mat = self.data_mat[:, 0:self.train_range[1]] self.scaler = self.scaler.fit(train_mat.transpose()) self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose() @@ -253,9 +248,8 @@ class TimeSeriesdata(object): gen_fn = self.train_gen else: gen_fn = lambda: self.test_val_gen(mode, shift) - output_types = tuple( - [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2 - ) + output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + + [tf.int32] * 2) dataset = tf.data.Dataset.from_generator(gen_fn, output_types) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) return dataset diff --git a/src/timesfm/timesfm.py b/src/timesfm/timesfm.py index ec63a20..f1f8b79 100644 --- a/src/timesfm/timesfm.py +++ b/src/timesfm/timesfm.py @@ -783,6 +783,7 @@ class TimesFm: model_name: str = "timesfm", window_size: int | None = None, num_jobs: int = 1, + verbose: bool = True, ) -> pd.DataFrame: """Forecasts on a list of time series. @@ -800,6 +801,7 @@ class TimesFm: window_size: window size of trend + residual decomposition. If None then we do not do decomposition. num_jobs: number of parallel processes to use for dataframe processing. + verbose: output model states in terminal. Returns: Future forecasts dataframe. @@ -819,7 +821,8 @@ class TimesFm: new_inputs = [] uids = [] if num_jobs == 1: - print("Processing dataframe with single process.") + if verbose: + print("Processing dataframe with single process.") for key, group in df_sorted.groupby("unique_id"): inp, uid = process_group( key, @@ -832,7 +835,8 @@ class TimesFm: else: if num_jobs == -1: num_jobs = multiprocessing.cpu_count() - print("Processing dataframe with multiple processes.") + if verbose: + print("Processing dataframe with multiple processes.") with multiprocessing.Pool(processes=num_jobs) as pool: results = pool.starmap( process_group, @@ -842,12 +846,14 @@ class TimesFm: ], ) new_inputs, uids = zip(*results) - print("Finished preprocessing dataframe.") + if verbose: + print("Finished preprocessing dataframe.") freq_inps = [freq_map(freq)] * len(new_inputs) _, full_forecast = self.forecast( new_inputs, freq=freq_inps, window_size=window_size ) - print("Finished forecasting.") + if verbose: + print("Finished forecasting.") fcst_df = make_future_dataframe( uids=uids, last_times=df_sorted.groupby("unique_id")["ds"].tail(1),