Update requirements and install instructions for multi-backend keras (k…

…eras-team#1257)

* Update requirements and install instructions for multi-backend keras

* Address comments

* less strict bounds

* Update to released tf-text

* Match keras versions

.github/workflows/actions.yml

@@ -30,7 +30,6 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -r requirements.txt --progress-bar off
-          pip install jax[cpu] --progress-bar off
           pip install --no-deps -e "." --progress-bar off
       - name: Test with pytest
         run: |
@@ -66,8 +65,6 @@ jobs:
     - name: Install dependencies
       run: |
           pip install -r requirements.txt --progress-bar off
-          pip install torch>=2.0.1+cpu --progress-bar off
-          pip install jax[cpu] --progress-bar off
           pip install --no-deps -e "." --progress-bar off
     - name: Test with pytest
       env:

CONTRIBUTING.md

@@ -81,7 +81,7 @@ Once the pull request is approved, a team member will take care of merging.
 
 ## Setting up an Environment
 
-Python 3.7 or later is required.
+Python 3.9 or later is required.
 
 Setting up your KerasNLP development environment requires you to fork the
 KerasNLP repository and clone it locally. With the
@@ -93,72 +93,57 @@ cd keras-nlp
 ```
 
 Next we must setup a python environment with the correct dependencies. We
-recommend using `conda` to install tensorflow dependencies (such as CUDA), and
-`pip` to install python packages from PyPI. The exact method will depend on your
-OS.
+recommend using `conda` to set up a base environment, and `pip` to install
+python packages from PyPI. The exact method will depend on your OS.
 
-**Note**: Please be careful not to use the `tensorflow` pre-packaged with conda,
-which is incompatible with `tensorflow-text` on PyPi, and follow the
-instructions below.
+**Note**: Be careful not to use mix pre-packaged tensorflow and jax libraries in
+`conda` with PyPI packages from `pip`. We recommend pulling *all* KerasNLP
+dependencies via `pip` as described below.
 
 ### Linux (recommended)
 
-To setup a complete environment with TensorFlow, a local install of keras-nlp,
-and all development tools, run the following or adapt it to suit your needs.
+For developing and unit testing the library, a CPU-only environment is often
+sufficient. For any training or inference with the library, you will quickly
+want accelerator support. The easiest way to get GPU support across all of our
+backends is to set up a few different python environements and pull in all cuda
+dependencies via `pip`.
+
+The shell snippet below will install four conda environments: `keras-nlp-cpu`,
+`keras-nlp-jax`, `keras-nlp-torch`, and `keras-nlp-tensorflow`. The cpu
+environement supports all backends without cuda, and each backend environement
+has cuda support.
 
 ```shell
-# Create and activate conda environment.
-conda create -n keras-nlp python=3.9
-conda activate keras-nlp
-
-# The following can be omitted if GPU support is not required.
-conda install -c conda-forge cudatoolkit-dev=11.2 cudnn=8.1.0
-mkdir -p $CONDA_PREFIX/etc/conda/activate.d/
-echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
-echo 'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
-source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
-
-# Install dependencies.
-python -m pip install --upgrade pip
-python -m pip install -r requirements.txt
-python -m pip install -e "."
+conda create -y -n keras-nlp-cpu python=3.10
+conda activate keras-nlp-cpu
+pip install -r requirements.txt  # install deps
+python pip_build.py --install  # install keras-nlp
+
+for backend in "jax" "torch" "tensorflow"; do
+    conda create -y -n keras-nlp-${backend} python=3.10
+    conda activate keras-nlp-${backend}
+    pip install -r requirements-${backend}-cuda.txt  # install deps
+    python pip_build.py --install  # install keras-nlp
+done
 ```
 
-### MacOS
-
-⚠️⚠️⚠️ MacOS binaries are for the M1 architecture are not currently available from
-official sources. You can try experimental development workflow leveraging the
-[tensorflow metal plugin](https://developer.apple.com/metal/tensorflow-plugin/)
-and a [community maintained build](https://github.com/sun1638650145/Libraries-and-Extensions-for-TensorFlow-for-Apple-Silicon)
-of `tensorflow-text`. These binaries are not provided by Google, so proceed at
-your own risk.
-
-#### Experimental instructions for Arm (M1)
+To activate the jax environment and set keras to use jax, run:
 
 ```shell
-# Create and activate conda environment.
-conda create -n keras-nlp python=3.9
-conda activate keras-nlp
-
-# Install dependencies.
-conda install -c apple tensorflow-deps=2.9
-python -m pip install --upgrade pip
-python -m pip install -r requirements-macos-m1.txt
-python -m pip install -e "."
+conda activate keras-nlp-jax && export KERAS_BACKEND=jax
 ```
 
-#### Instructions for x86 (Intel)
+### MacOS
 
-```shell
-# Create and activate conda environment.
-conda create -n keras-nlp python=3.9
-conda activate keras-nlp
-
-# Install dependencies.
-python -m pip install --upgrade pip
-python -m pip install -r requirements.txt
-python -m pip install -e "."
-```
+`tensorflow-text` does not release precompiled binaries for MacOS M-series
+chips, though the library does support building from source on MacOS.
+
+We strongly recommend a Linux development environment for an easy contribution
+experience. To build a dev environement from scratch on MacOS, see the following
+guides:
+
+https://developer.apple.com/metal/tensorflow-plugin/
+https://github.com/tensorflow/text
 
 ### Windows
 

requirements-common.txt

@@ -16,3 +16,6 @@ namex
 # Optional deps.
 rouge-score
 sentencepiece
+tensorflow-datasets
+# Breakage fix.
+ml-dtypes==0.2.0

requirements-jax-cuda.txt

@@ -0,0 +1,14 @@
+# Tensorflow cpu-only version.
+tensorflow>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch cpu-only version.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax with cuda support.
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda12_pip]
+
+-r requirements-common.txt

requirements-tensorflow-cuda.txt

@@ -0,0 +1,13 @@
+# Tensorflow with cuda support.
+tensorflow[and-cuda]>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch cpu-only version.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax cpu-only version.
+jax[cpu]
+
+-r requirements-common.txt

requirements-torch-cuda.txt

@@ -0,0 +1,13 @@
+# Tensorflow cpu-only version.
+tensorflow>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch with cuda support.
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.1.0
+torchvision==0.16.0
+
+# Jax cpu-only version.
+jax[cpu]
+
+-r requirements-common.txt

requirements.txt

@@ -1,7 +1,13 @@
-# Core deps.
-tensorflow~=2.13.0
-tensorflow-text~=2.13.0
-tensorflow-datasets
+# Tensorflow.
+tensorflow>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax.
+jax[cpu]
 
-# Common deps.
 -r requirements-common.txt