Merge branch 'master' of github.com:google-research/timesfm

2024-07-12 22:18:25 +01:00
parent d277f3bd69 0a40da3c52
commit 203370d1a5
5 changed files with 231 additions and 25 deletions
@@ -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.
@@ -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"
@@ -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"
@@ -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
@@ -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),