support conda env

2024-12-26 19:49:54 +03:00 · 2024-03-20 11:19:03 +08:00 · 2024-03-20 11:19:03 +08:00 · f78574c4eb
commit f78574c4eb
parent 7050ed204b
2 changed files with 19 additions and 0 deletions
--- a/README.md
+++ b/README.md
@ -11,6 +11,14 @@ pip install -r requirements.txt
 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.
--- a/environment.yml
+++ 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