support conda env

This commit is contained in:
huangganggui 2024-03-20 11:19:03 +08:00
parent 7050ed204b
commit f78574c4eb
2 changed files with 19 additions and 0 deletions

View File

@ -11,6 +11,14 @@ pip install -r requirements.txt
python run.py 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. to test the code.
The script loads the checkpoint and samples from the model on a test input. The script loads the checkpoint and samples from the model on a test input.

11
environment.yml Normal file
View File

@ -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