From bd57aeda9e3e13ed514f08d25ad62b9de951642c Mon Sep 17 00:00:00 2001
From: David Landup <60978046+DavidLandup0@users.noreply.github.com>
Date: Wed, 30 Oct 2024 11:25:56 +0900
Subject: [PATCH] [T5 1.1] Enable v1.1 Presets (#1948)

* enable t5 1.1 weights

* add xxl config

* consolidate usage of relu vs geglu vs variable activation functions

* update preset param counts and xll config

* remove commented code

* revert to use_gated_activation

* add kaggle links

* fix comment

* remove xl preset

* update kerashub links to keras links
---
 keras_hub/.DS_Store                           | Bin 0 -> 6148 bytes
 keras_hub/src/models/t5/t5_backbone.py        |   9 +-
 keras_hub/src/models/t5/t5_presets.py         |  32 ++++-
 tools/.DS_Store                               | Bin 0 -> 8196 bytes
 .../convert_t5_checkpoints.py                 | 110 +++++++++++++++---
 5 files changed, 128 insertions(+), 23 deletions(-)
 create mode 100644 keras_hub/.DS_Store
 create mode 100644 tools/.DS_Store

diff --git a/keras_hub/.DS_Store b/keras_hub/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..4147d62e9eab062b0a418ce3f8fa741f925aacf5
GIT binary patch
literal 6148
zcmeHKu};G<5IrXyipo$n#_$RKfJ&hX3u_tJTACt?)G7glRG9J$EUXMnd;u#HpTx|&
zJB!2$MPdLUbSK#_Ip5_u&yg=CA~(984~Y6i)J0>Aw$OcIJkG6Uk@su^ojfC_6Uu2p
zr!;9ro8vDkz-!mTYE7u58?0Tuev{?dV4N4LQl}mFHb<_CESoKh86vueN5lKO<JbMK
zqUYa;E*H&N@OMfvW?9mCIKStodZam@ov3)Yn8|#6-m0ki+oJ)Wn=0QGT~Z$K-M$Tl
zPez)2>%Dc&=SOS4_GrNO_`-Z<rhqA63jDMJsM%~?-r3De0aL&fXez+_LxRSbdaM*>
z>p&+*0AL4hF<i?(8Zco8F!fj|!UJ=f3e;5Rj~GtV;ZJ;A>akMPbaKLs=Xhr4Pbf~<
z;ZG!-T&ifjDPRhO3dDNb=lXxJ`TQRy*_|n13j8YtTqhf6BfOHYt*w{iTAQFv(bzby
mQv684g_L5%aw*<L6N5kD3t;N8QiKO)KLR3yb*8|tD)0`=&}IGr

literal 0
HcmV?d00001

diff --git a/keras_hub/src/models/t5/t5_backbone.py b/keras_hub/src/models/t5/t5_backbone.py
index 7e351cba62..103cf253e5 100644
--- a/keras_hub/src/models/t5/t5_backbone.py
+++ b/keras_hub/src/models/t5/t5_backbone.py
@@ -42,11 +42,12 @@ class T5Backbone(Backbone):
             projections in the multi-head attention layers. Defaults to
             hidden_dim / num_heads.
         dropout: float. Dropout probability for the Transformer layers.
-        activation: activation function (or activation string name). The
-            activation to be used in the inner dense blocks of the
-            Transformer layers. Defaults to `"relu"`.
+        activation: string. The activation function to use in the dense blocks
+            of the Transformer Layers.
         use_gated_activation: boolean. Whether to use activation gating in
-            the inner dense blocks of the Transformer layers.
+            the inner dense blocks of the Transformer layers. When used with
+            the GELU activation function, this is referred to as GEGLU
+            (gated GLU) from https://arxiv.org/pdf/2002.05202.
             The original T5 architecture didn't use gating, but more
             recent versions do. Defaults to `True`.
         layer_norm_epsilon: float. Epsilon factor to be used in the
diff --git a/keras_hub/src/models/t5/t5_presets.py b/keras_hub/src/models/t5/t5_presets.py
index 72edf24b4c..109f6d6981 100644
--- a/keras_hub/src/models/t5/t5_presets.py
+++ b/keras_hub/src/models/t5/t5_presets.py
@@ -1,4 +1,4 @@
-"""XLM-RoBERTa model preset configurations."""
+"""T5 model preset configurations."""
 
 backbone_presets = {
     "t5_small_multi": {
@@ -14,6 +14,16 @@
         },
         "kaggle_handle": "kaggle://keras/t5/keras/t5_small_multi/2",
     },
+    "t5_1.1_small": {
+        "metadata": {
+            "description": (""),
+            "params": 60511616,
+            "official_name": "T5 1.1",
+            "path": "t5",
+            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_small",
+    },
     "t5_base_multi": {
         "metadata": {
             "description": (
@@ -27,6 +37,16 @@
         },
         "kaggle_handle": "kaggle://keras/t5/keras/t5_base_multi/2",
     },
+    "t5_1.1_base": {
+        "metadata": {
+            "description": (""),
+            "params": 247577856,
+            "official_name": "T5 1.1",
+            "path": "t5",
+            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_base",
+    },
     "t5_large_multi": {
         "metadata": {
             "description": (
@@ -40,6 +60,16 @@
         },
         "kaggle_handle": "kaggle://keras/t5/keras/t5_large_multi/2",
     },
+    "t5_1.1_large": {
+        "metadata": {
+            "description": (""),
+            "params": 750251008,
+            "official_name": "T5 1.1",
+            "path": "t5",
+            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_large",
+    },
     "flan_small_multi": {
         "metadata": {
             "description": (
diff --git a/tools/.DS_Store b/tools/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..b8a3ee4883e2e53c53215edceedaeb67daff4ec6
GIT binary patch
literal 8196
zcmeHMu}&L75S=BqWXUNa1teGKDAOddXd`hTr%Hq5A|+t^WKeA10fvI=L_wYW0pEZq
zD5+9LY4`zJeu6hM>-a9VB_Bu@60_FMySv@l+4sD@y*U$++dApgiE2btL1lbbg_%;g
zo$Ewd@>iCD4*o<xD501RXh6e>-gdzhFa=BjQ@|831zv&z*t5A*X6*agYSt7m1zx2B
ze18Z~8AFe$qTD)QBn1GLa9IkDkq3xR;4$==DvB${)a=1|SLag<r|xJExH~ZPm@29}
zIdv!JBRij=IEfCwK)RC)70sFgrof~EoV%}3oAxP!&F1erIFFIZpQ7c6j*->ltk<nK
z+i^dUfl7lf#k~7*qtWWcEqGO~s^9OH+rKx9j?j7L2+4kiLHQ>Wc|U{y5WM%~Z3e&5
z=;pZ7`o;X7S>E>_um9X^3-3AO_Tw?<;@&yFdNVCk7vE_gpS?V%ix1$GY0b|2*Wb0v
z!<NiD*JC`l={@rCFA)IiErGHT-Vu%Tji^IAIFGRcBWQg<XS(Y0^LzVH%sV^(^KYx8
z`4@<v)&6NC<oSPA{>Q>{D7FoofHD96B2gj2uWcyn?!Prx`y~A5_RjL(+z7_}3yd)S
zexEi%p8u^rO|aDs-1dK75mror#VhbeL)JL||2q5q|KcZR5lsP8U||ZlQlr_}LO7pW
yPXy+zIzZh;<-+|^MR9?V-5e<CIPl*ehByw;m3=~wsUohR*n@z`V8#?!r~;39V+nHr

literal 0
HcmV?d00001

diff --git a/tools/checkpoint_conversion/convert_t5_checkpoints.py b/tools/checkpoint_conversion/convert_t5_checkpoints.py
index e0717174ad..fc33919410 100644
--- a/tools/checkpoint_conversion/convert_t5_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_t5_checkpoints.py
@@ -7,7 +7,6 @@
 from absl import app
 from absl import flags
 from checkpoint_conversion_utils import get_md5_checksum
-from keras import ops
 
 import keras_hub
 
@@ -15,11 +14,91 @@
     "t5_small_multi": "t5-small",
     "t5_base_multi": "t5-base",
     "t5_large_multi": "t5-large",
+    "t5_1.1_small": "google/t5-v1_1-small",
+    "t5_1.1_base": "google/t5-v1_1-base",
+    "t5_1.1_large": "google/t5-v1_1-large",
+    "t5_1.1_xl": "google/t5-v1_1-xl",
+    "t5_1.1_xxl": "google/t5-v1_1-xxl",
     "flan_small_multi": "google/flan-t5-small",
     "flan_base_multi": "google/flan-t5-base",
     "flan_large_multi": "google/flan-t5-large",
 }
 
+
+PARAM_MAP = {
+    "t5_1.1_small": {
+        "trainable": True,
+        "vocabulary_size": 32128,
+        "hidden_dim": 512,
+        "intermediate_dim": 1024,
+        "num_layers": 8,
+        "num_heads": 6,
+        "activation": "gelu",
+        "key_value_dim": 64,
+        "dropout": 0.1,
+        "use_gated_activation": True,
+        "layer_norm_epsilon": 1e-6,
+        "tie_embedding_weights": False,
+    },
+    "t5_1.1_base": {
+        "trainable": True,
+        "vocabulary_size": 32128,
+        "hidden_dim": 768,
+        "intermediate_dim": 2048,
+        "num_layers": 12,
+        "num_heads": 12,
+        "activation": "gelu",
+        "key_value_dim": 64,
+        "dropout": 0.1,
+        "use_gated_activation": True,
+        "layer_norm_epsilon": 1e-6,
+        "tie_embedding_weights": False,
+    },
+    "t5_1.1_large": {
+        "trainable": True,
+        "vocabulary_size": 32128,
+        "hidden_dim": 1024,
+        "intermediate_dim": 2816,
+        "num_layers": 24,
+        "num_heads": 16,
+        "activation": "gelu",
+        "key_value_dim": 64,
+        "dropout": 0.1,
+        "use_gated_activation": True,
+        "layer_norm_epsilon": 1e-6,
+        "tie_embedding_weights": False,
+    },
+    "t5_1.1_xl": {
+        "trainable": True,
+        "vocabulary_size": 32128,
+        "hidden_dim": 2048,
+        "intermediate_dim": 5120,
+        "num_layers": 24,
+        "num_heads": 32,
+        "activation": "gelu",
+        "key_value_dim": 64,
+        "dropout": 0.1,
+        "use_gated_activation": True,
+        "layer_norm_epsilon": 1e-6,
+        "tie_embedding_weights": False,
+    },
+    "t5_1.1_xxl": {
+        "trainable": True,
+        "vocabulary_size": 32128,
+        "hidden_dim": 4096,
+        "intermediate_dim": 10240,
+        "num_layers": 24,
+        "num_heads": 64,
+        "activation": "gelu",
+        "key_value_dim": 64,
+        "dropout": 0.1,
+        "use_gated_activation": True,
+        "layer_norm_epsilon": 1e-6,
+        "tie_embedding_weights": False,
+    },
+}
+
+
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string(
@@ -52,9 +131,7 @@ def extract_vocab(hf_tokenizer):
 
 
 def convert_checkpoints(hf_model):
-    keras_hub_model = keras_hub.models.T5Backbone.from_preset(
-        FLAGS.preset, load_weights=False
-    )
+    keras_hub_model = keras_hub.models.T5Backbone(**PARAM_MAP[FLAGS.preset])
 
     hf_wts = hf_model.state_dict()
     print("Original weights:")
@@ -308,17 +385,12 @@ def check_output(
     keras_hidden_states = keras_out["decoder_sequence_output"]
     hf_hidden_states = hf_out.decoder_hidden_states[-1]
 
-    keras_outputs = ops.take_along_axis(
-        keras_hidden_states, ops.where(decoder_padding_mask)
-    )
-    hf_outputs = ops.take_along_axis(
-        hf_hidden_states, ops.where(decoder_padding_mask)
-    )
-
-    print("-> KerasHub output:", keras_outputs[0:5])
-    print("-> HF output:", hf_outputs[0:5])
+    print("-> KerasHub output:", keras_hidden_states[0:5])
+    print("-> HF output:", hf_hidden_states[0:5])
     np.testing.assert_allclose(
-        keras_outputs.detach().numpy(), hf_outputs.detach().numpy(), atol=1e-5
+        keras_hidden_states.numpy(),
+        hf_hidden_states.detach().numpy(),
+        atol=1e-2,
     )
 
     if keras_model.tie_embedding_weights:
@@ -333,7 +405,7 @@ def check_output(
     print("-> KerasHub logits:", keras_logits[0:5])
     print("-> HF logits:", hf_logits[0:5])
     np.testing.assert_allclose(
-        keras_logits.detach().numpy(), hf_logits.detach().numpy(), atol=1e-3
+        keras_logits.numpy(), hf_logits.detach().numpy(), atol=1e-1
     )
 
 
@@ -352,16 +424,18 @@ def main(_):
     keras_model = convert_checkpoints(hf_model)
 
     # Save the model.
-    model_path = f"./{FLAGS.preset}/model.weights.h5"
+    model_path = f"./{FLAGS.preset}"
+    weight_path = os.path.join(model_path, "model.weights.h5")
     print(f"\n-> Save KerasHub model weights to `{model_path}`.")
-    keras_model.save_weights(model_path)
+    keras_model.save_to_preset(model_path)
     print("-> Print MD5 checksum of the model weights files.")
-    print(f"`{model_path}` md5sum: ", get_md5_checksum(model_path))
+    print(f"`{model_path}` md5sum: ", get_md5_checksum(weight_path))
     print(f"-> Param count {count_params(keras_model.weights)}")
 
     print("\n-> Convert vocab.")
     hf_tokenizer = transformers.AutoTokenizer.from_pretrained(hf_id)
     keras_tokenizer = extract_vocab(hf_tokenizer)
+    keras_tokenizer.save_to_preset(model_path)
 
     check_output(
         keras_model,