I've been working on a personal project called νοῦς (Nous). It's a Python implementation of the transformer model described in Attention is All You Need (Vaswani et al., 2017), using only JAX for the backend. The goal started as a learning exercise, but the system turned into something that many others might also find useful.

As I previously mentioned, the backend is completely Python, but I have also built an electron-based frontend, releasing a fully built MacOS app. If you are on Windows, you can clone the repository and visit the README for further instructions. All model components are completely configurable: depth, width, number of heads, training parameters, and generation settings. You can both train locally using the Electron app or using the CLI with a remote GPU. Note that using the CLI, you will have to modify source code to edit the components.

Nous would be useful in experimenting with specific architectures, running controlled training setups, or even teaching/learning how modern LLMs work (the README also explains how every single part of my backend works, so it could be quite helpful to use while reading the codebase).

I originally wanted to make it a NumPy-only version but switched fully to JAX after running into performance problems.

Currently, the app comes with a pre-trained model that has 76.9M parameters, trained to a final loss value of ~0.6. It also comes with a Byte-Pair Encoding tokenizer implementation, so that you can build the tokenizer on a dataset of your choice.

Everything is open-source. Repo link is above.