Support for cpu/gpu version

.github/workflows/main.yml

@@ -20,7 +20,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install poetry
-      # use the hardcoded version number for now
+      # use the hardcoded version number for now. if this is merged to the main repo, you can uncomment these lines to automatically tag the version number with the github run number
       # - name: Set Version number
       #   run: |
       #     poetry version 0.0.${{ github.run_number }}

.gitignore

@@ -1,2 +1,3 @@
 .venv/
 dist/
+**__pycache__/** */

README.md

@@ -14,6 +14,8 @@ to download model checkpoints.
 
 This is not an officially supported Google product.
 
+We recommend at least 16GB RAM to load TimesFM dependencies.
+
 ## Checkpoint timesfm-1.0-200m
 
 timesfm-1.0-200m is the first open model checkpoint:
@@ -32,51 +34,28 @@ Please look into the README files in the respective benchmark directories within
 
 ### Installation as a package
 
-To install the CPU version of TimesFM as a package, you can run the following command without cloning this repo:
+To install the GPU/default version of TimesFM as a package, you can run the following command without cloning this repo:
 
 `pip install timesfm`
 
 The package is currently using https://pypi.org/project/timesfm/.
 
-### Poetry (CPU) installation
+To install the CPU version of TimesFM, you can run the command `pip install timesfm[cpu]`
 
-To install the CPU environment using Poetry, you can run the following command:
+### Local installation
+
+To from the current repository/local version (like you would have previously done with `pip -e .`), you can run the command
 
 ```
 pip install poetry # optional
 poetry install
 ```
 
+If you prefer the CPU version, run the command `poetry install -E cpu`
+
 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.
 
-### Conda / GPU installation
+## Notes
-
-We recommend at least 16GB RAM to load TimesFM dependencies.
-
-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.

environment.yml

@@ -1,18 +0,0 @@
-name: tfm_env
-
-channels:
-  - conda-forge
-  - defaults
-  - anaconda
-dependencies:
-  - jupyterlab
-  - pip
-  - python=3.10
-  - pip:
-    - huggingface_hub[cli]
-    - utilsforecast
-    - praxis
-    - paxml
-    - jax[cuda12]==0.4.26
-    - einshape
-    - scikit-learn

environment_cpu.yml

@@ -1,18 +0,0 @@
-name: tfm_env
-
-channels:
-  - conda-forge
-  - defaults
-  - anaconda
-dependencies:
-  - jupyterlab
-  - pip
-  - python=3.10
-  - pip:
-    - huggingface_hub[cli]
-    - utilsforecast
-    - praxis
-    - paxml
-    - jax[cpu]==0.4.26
-    - einshape
-    - scikit-learn

pyproject.toml

@@ -1,25 +1,10 @@
-# This project can be installed with `python3 -m pip install -e .` from the main directory. 
+[tool.poetry]
-
+name = "timesfm"
-[project]
-name = "timesfm-jax"
 packages = [
     { include = "*", from = "src" },
 ]
-version = "0.0.1"
 description = "Open weights time-series foundation model from Google Research."
-version = "1.0.1"
+version = "0.0.12"
-dependencies = [
-    "jax==0.4.26",
-    "paxml==1.4.0",
-    "praxis==1.4.0",
-    "jaxlib==0.4.26",
-    "numpy==1.26.4",
-    "pandas==2.1.4",
-    "einshape==1.0.0",
-    "utilsforecast==0.1.10",
-    "huggingface_hub[cli]==0.23.0", 
-    "scikit-learn==1.5.1",
-]
 authors = [
     "Rajat Sen <senrajat@google.com>",
     "Yichen Zhou <yichenzhou@google.com>",
@@ -49,9 +34,13 @@ einshape = ">=1.0.0"
 numpy = ">=1.26.4"
 pandas = ">=2.1.4"
 paxml = "1.4.0"
-jax = "0.4.26"
-jaxlib = "0.4.26"
 utilsforecast = "^0.1.10"
+jax = {version = "0.4.26", extras = ["cuda12"]}
+# unfortunately we have to solve it using https://stackoverflow.com/questions/72037181/how-to-add-optional-dependencies-of-a-library-as-extra-in-poetry-and-pyproject
+Jax = {version = "0.4.26", extras = ["cpu"], optional = true}
+
+[tool.poetry.extras]
+cpu = ['Jax']
 
 
 [build-system]