setup_tpu_vm_venv.sh

# setup training on a TPU VM
rm -fr venv
python3 -m venv venv
source venv/bin/activate
pip install -U pip
pip install -U wheel
pip install requests
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

cd ..
git clone https://github.com/huggingface/transformers.git
cd transformers
pip install -e ".[flax]"
cd ..

git clone https://github.com/huggingface/datasets.git
cd datasets
pip install -e ".[streaming]"
cd ..