Merge c3f5f2dc12 into 7050ed204b

README.md

@@ -13,6 +13,16 @@ python run.py
 
 to test the code.
 
+If you are using Conda, run
+
+```shell
+conda env create -f environment.yml
+conda activate grok
+python run.py
+```
+
+to test the code.
+
 The script loads the checkpoint and samples from the model on a test input.
 
 Due to the large size of the model (314B parameters), a machine with enough GPU memory is required to test the model with the example code.

environment.yml

@@ -0,0 +1,11 @@
+name: grok
+channels:
+  - conda-forge
+  - nvidia
+  - defaults
+dependencies:
+  - python=3.9
+  - numpy==1.26.4
+  - dm-haiku==0.0.12
+  - sentencepiece==0.2.0
+  - jax[cuda12-pip]==0.4.25