kira_llm 0.1.1

No project description provided

That's right. I'm Kira ✍️

Kira is a transformer built with JAX and Equinox, which is IMHO the best neural network library built on top of JAX out there, because of its simple and elegant design.

The role of Kira is to serve as a baseline transformer implementation, which is very easy to understand and extend.

Kira itself does not yet allow for KV caching (although there is a branch in which I'm trying to get it to work). It does however allow you to interpolate between Multi-Query Attention (MQA) and "regular" Multi-Head Attention (MHA). This is different from other MHA implementations, where most of the time you can only set the number of heads and that's it. Kira offers more flexibility in that regard.

These features of Kira (the interpolation between MQA and MHA and the RoPE embeddings) will soon be integrated in the main Equinox repository at which point Kira's MHA implementation will be replaced with the built-in Equinox's MHA.

---

To get started with Kira, you can either install it with

pip3 install kira_llm

or simply clone the repository and cherry-pick what you need.

Contributing

To contribute to this project, you'll need to fork this repository and run

poetry install

which will install all of the dependencies. Then, simply make your changes and start a pull request. The philosophy of this repository is to be simple and understandable.