Upload 2 files

colab_env.txt

@@ -0,0 +1,23 @@
+airio @ git+https://github.com/google/airio#egg=airio
+absl-py==1.4.0
+clu @ git+https://github.com/google/CommonLoopUtils#egg=clu
+flax @ git+https://github.com/google/flax#egg=flax
+fiddle==0.2.11
+gin-config==0.5.0
+jax==0.4.20
+jaxtyping==0.2.24
+jaxlib @ https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.20+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl#sha256=01be66238133f884bf5adf15cd7eaaf8445f9d4b056c5c64df28a997a6aff2fe
+jestimator @ git+https://github.com/google-research/jestimator#egg=jestimator
+numpy==1.23.5
+optax @ git+https://github.com/deepmind/optax@1a7956d0e6f6f2881fb8f7d4c579a6435c8431eb
+orbax-checkpoint==0.4.4
+seqio @ git+https://github.com/google/seqio#egg=seqio
+tensorflow-cpu==2.15.0.post1
+tensorstore==0.1.45
+protobuf==3.20.3
+colossalai >= 0.3.3
+torch >= 1.8.1
+transformers==4.34.0
+sentencepiece==0.1.99
+datasets==2.14.7
+chex==0.1.7

prepare_env.sh

@@ -0,0 +1,7 @@
+python --version
+python -m pip install --upgrade pip
+
+pip install -r colab_env.txt
+# Install t5x
+git clone --branch=main https://github.com/google-research/t5x
+python -m pip install ./t5x