kira_llm 0.1.2

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 now allows for KV caching. It also allows 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.

Kira can also be used as an encoder. Simply pass mask=None when you call Kira and there will be no masking in the MHA (i.e. making it an encoder).

---

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.