This summer, I spent my days (and too many late nights) turning Pomni from a rough proposal into something real. I joined Google Summer of Code with the goal of tackling a simple but pressing problem: the lack of AI tooling that integrates cleanly with the JAX/Keras ecosystem. Over a few months, I made significant progress toward closing that gap.

Current frontier LLMs struggle with JAX. Resources are fragmented, documentation is thin, and examples are scattered across different platforms. As a result, even basic assistance—like debugging or writing idiomatic JAX code—remains a pain point for developers.

That’s where https://github.com/neel04/pomni comes in. It’s an end-to-end pipeline for fine-tuning off-the-shelf Large Language Models, with a particular focus on the open-source Gemma series. It also includes a lightweight terminal UI so developers can integrate it directly into their workflows without extra overhead.

The process came with challenges. A major issue was that Gemma often lost its RLHF-trained personality after fine-tuning, slipping into a more pedantic and sometimes abrasive tone—likely influenced by the large amount of Stack Overflow data in the training set. Preserving helpfulness while retaining accuracy required careful iteration.

Resource constraints were another limiting factor. Training and testing larger models directly wasn’t feasible, so I relied on efficiency tricks—primarily quantization and lower-precision training—to make smaller-scale experiments practical. These methods allowed me to validate ideas and push the models further than raw compute would have otherwise allowed.

For evaluation, I used a smaller version of Gemini 2.0 as an automatic judge, comparing outputs between base models and their Pomni-tuned counterparts. This provided consistent, structured feedback, which proved invaluable in guiding fine-tuning toward models that were not only more capable, but also more cooperative.

Keras was central to making all of this work. Its streamlined support for training and inference, combined with the JAX backend, allowed the entire pipeline to remain efficient, modular, and easy to extend.

Looking forward, I hope to extend Pomni to fine-tuning larger Gemini models. This would make them significantly more capable for practical use, while also more approachable and useful for developers working within the JAX/Keras ecosystem.