Add T5Gemma to KerasHub

[harshaljanjani]

Description of the change

Closes the issue #2321

Numerics Consistency Check and Example Output

from keras_hub.src.models.t5gemma.t5gemma_causal_lm import T5GemmaCausalLM

t5gemma_lm = T5GemmaCausalLM.from_preset("t5gemma-b-b-prefixlm-it")
output = t5gemma_lm.generate("What is the fastest land animal?")
print(output)

Reference

• Encoder-Decoder Gemma: Improving the Quality-Efficiency Trade-Off via Adaptation
• Pre-trained weights from Hugging Face

Colab Notebook

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have followed the Keras Hub Model contribution guidelines in making these changes.
• I have followed the Keras Hub API design guidelines in making these changes.
• I have signed the Contributor License Agreement.

[gemini-code-assist]

Summary of Changes

Hello @harshaljanjani, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates the T5Gemma model into KerasHub, providing a comprehensive implementation of its encoder-decoder architecture, attention mechanisms, and supporting components. It enables causal language modeling with efficient text generation capabilities and includes dedicated preprocessing and tokenization utilities.

Highlights

• New Model Integration: This pull request introduces the complete T5Gemma model architecture into KerasHub, enabling its use for various natural language processing tasks.
• Advanced Attention Mechanisms: New T5GemmaSelfAttention and T5GemmaCrossAttention layers are added, featuring support for Grouped Query Attention (GQA) and Rotary Positional Embeddings (RoPE) for enhanced performance and positional encoding.
• Encoder-Decoder Backbone: The T5GemmaBackbone is implemented, providing the core encoder-decoder structure. It supports both full attention and sliding window attention within its layers.
• Causal Language Modeling: An end-to-end T5GemmaCausalLM is included, designed for efficient text generation through optimized call_with_cache and generate_step methods for autoregressive inference.
• Dedicated Preprocessing and Tokenization: Custom T5GemmaCausalLMPreprocessor and T5GemmaTokenizer classes are added to handle input data preparation, including tokenization with SentencePiece and management of special tokens.
• Core Layer Components: Fundamental building blocks like T5GemmaMLP (Multi-Layer Perceptron) and a specific t5gemma_kernel_initializer are introduced to support the T5Gemma architecture.
• Comprehensive Testing: New unit tests are provided for the T5GemmaBackbone and T5GemmaCausalLM to ensure the correctness of the implementation and proper model saving functionality.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The code introduces the T5Gemma model to KerasHub. The implementation is comprehensive and well-structured. The review focuses on a performance optimization for the generation process and a point of code consistency. Addressing these will enhance the model's efficiency and maintainability.