Add T5Gemma to KerasHub (#2339)

* init: Add initial project structure and files

* nit: Fix code format test; and cool AI-generated reviews

* refactor: Cleanup and replace incorrect T5LayerNorm with RMSNormalization (Gemma)

* fix: Numerics @ atol=1e-4

* refactor: Refactor T5Gemma decoder cache handling

* feat: Add checkpoint conversion script

* nit: Precise compute_output_shape methods; document head_dim

* nit: Propagate dtypes

* bug fix + minor cleanup: Fix head_dim default → head_dim from config

* perf(jax/tpu): Fused kernel optim for TPU backend + get_config() args

* cleanup: Slight refactor

* fix: Enable mixed precision and quantization tests

* feat: Add support for asymmetrical presets (only invariants included)

* refactor: Address reviews - presets will be handled post D-FINE

* feat: Support direct loading of Hugging Face checkpoints

* ✅ Yayy: Generate outputs identical, hidden states match within 1e-3

* preset test: Register and test a preset (to be replaced later by the team with the full set)

* nit: Sharded weights don’t include `model.weights.h5`

* nits: Address reviews + replace gated model

Author: harshaljanjani

keras_hub/api/models/__init__.py

@@ -615,6 +615,18 @@
     T5Preprocessor as T5Preprocessor,
 )
 from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer as T5Tokenizer
+from keras_hub.src.models.t5gemma.t5gemma_backbone import (
+    T5GemmaBackbone as T5GemmaBackbone,
+)
+from keras_hub.src.models.t5gemma.t5gemma_seq_2_seq_lm import (
+    T5GemmaSeq2SeqLM as T5GemmaSeq2SeqLM,
+)
+from keras_hub.src.models.t5gemma.t5gemma_seq_2_seq_lm_preprocessor import (
+    T5GemmaSeq2SeqLMPreprocessor as T5GemmaSeq2SeqLMPreprocessor,
+)
+from keras_hub.src.models.t5gemma.t5gemma_tokenizer import (
+    T5GemmaTokenizer as T5GemmaTokenizer,
+)
 from keras_hub.src.models.task import Task as Task
 from keras_hub.src.models.text_classifier import TextClassifier as Classifier
 from keras_hub.src.models.text_classifier import (

keras_hub/api/tokenizers/__init__.py

@@ -88,6 +88,9 @@
     SigLIPTokenizer as SigLIPTokenizer,
 )
 from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer as T5Tokenizer
+from keras_hub.src.models.t5gemma.t5gemma_tokenizer import (
+    T5GemmaTokenizer as T5GemmaTokenizer,
+)
 from keras_hub.src.models.whisper.whisper_tokenizer import (
     WhisperTokenizer as WhisperTokenizer,
 )

keras_hub/src/models/t5gemma/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.t5gemma.t5gemma_backbone import T5GemmaBackbone
+from keras_hub.src.models.t5gemma.t5gemma_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, T5GemmaBackbone)