This repository is based on nanoGPT from karpathy and you will find:
nanoGPT.ipynb, the original notebook in PyTorchnanoGPT_jax.ipynb, the original notebook translated in JAXnanoGPT_jax.py, a script to train a nanoGPT of ~200 lines in JAX
You can find an example of the text it generates after training it for ~30 minutes in outputs/output.txt.
To run this code, you need to clone the repository and install the required libraries with:
git clone ...
pip install -r requirements.txtTo build a container using Apptainer make sure you are in the root of the repository and then run:
apptainer build --fakeroot --force apptainer/container.sif apptainer/container.def