From f78574c4ebfbf38667d33b3ab43f2b6cf5aa6120 Mon Sep 17 00:00:00 2001
From: huangganggui <huangganggui@gmail.com>
Date: Wed, 20 Mar 2024 11:19:03 +0800
Subject: [PATCH] support conda env

---
 README.md       |  8 ++++++++
 environment.yml | 11 +++++++++++
 2 files changed, 19 insertions(+)
 create mode 100644 environment.yml

diff --git a/README.md b/README.md
index f501a07..26af219 100644
--- 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.
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