Merge pull request #2 from pcuenca/main

README.md

@@ -33,6 +33,10 @@
 
 
 ## Dependencies Installation
-You should create a new python virtual environment and install the project dependencies inside the virtual env: `pip install -r requirements.txt`
+You should create a new python virtual environment and install the project dependencies inside the virtual env. You need to use the `-f` (`--find-links`) option for `pip` to be able to find the appropriate `libtpu` required for the TPU hardware:
+
+```
+$ pip install -r requirements.txt -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+```
 
 If you use `conda`, you can create the virtual env and install everything using: `conda env update -f environments.yaml`

requirements.txt

@@ -1,5 +1,9 @@
-jax
+# Note: install with the following command:
+# pip install -r requirements.txt -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+# Otherwise it won't find the appropriate libtpu_nightly
+requests
+jax[tpu]>=0.2.16
+-e git+https://github.com/huggingface/transformers.git@master#egg=transformers
+-e git+https://github.com/huggingface/datasets.git@master#egg=datasets
 flax
-optax
-datasets
-git+https://github.com/huggingface/transformers
+jupyter