diff --git a/README.md b/README.md index f501a07..aa6396d 100644 --- a/README.md +++ b/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. diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..825f8ee --- /dev/null +++ b/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