Add T5LA models by HRezaei · Pull Request #39293 · huggingface/transformers

HRezaei · 2025-07-09T06:43:41Z

What does this PR do?

Adds the implementation of the LookAhead (LA) models. These models are designed to predict not only the next immediate token after the input prompt, but also the second, third, ... up to K next tokens. These tokens can be used to mitigate the high inference latency in generation (see this survey) or in approximated ranking of a set of responses (see this paper for an application).

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

@ArthurZucker

HRezaei · 2025-07-09T06:54:55Z

        ("swinv2", "Swin Transformer V2"),
        ("switch_transformers", "SwitchTransformers"),
        ("t5", "T5"),
+        ("t5la", "T5LALA"),


Why add-new-model-like has added two "LA" suffixes? It has to be "T5LA" I guess. I double checked my answers to that command and even retried it a second time. Everything was correct but this output doesn't look right!

ArthurZucker

Thanks for adding the model! 🚀

ArthurZucker · 2025-07-15T14:03:36Z

        ("switch_transformers", "SwitchTransformers"),
        ("t5", "T5"),
        ("t5gemma", "T5Gemma"),
+        ("t5la", "T5LALA"),


A bit wird to have this no?

ArthurZucker · 2025-07-15T14:03:57Z

+from collections.abc import Mapping
+
+from ...configuration_utils import PretrainedConfig
+from ...onnx import OnnxSeq2SeqConfigWithPast


no longer used for new models!

ArthurZucker · 2025-07-15T14:04:19Z

+    """
+
+    model_type = "t5la"
+    lookahead_type = "la"  # Other options are "laa", "laa2", and "lae"


think we should add this as an arg no?

ArthurZucker · 2025-07-15T14:04:40Z

+        # for backwards compatibility
+        if feed_forward_proj == "gated-gelu":
+            self.dense_act_fn = "gelu_new"


new model don't have BC

ArthurZucker · 2025-07-15T14:04:49Z

+class T5LAOnnxConfig(OnnxSeq2SeqConfigWithPast):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        common_inputs = {
+            "input_ids": {0: "batch", 1: "encoder_sequence"},
+            "attention_mask": {0: "batch", 1: "encoder_sequence"},
+        }
+        if self.use_past:
+            common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence"
+            common_inputs["decoder_input_ids"] = {0: "batch"}
+            common_inputs["decoder_attention_mask"] = {0: "batch", 1: "past_decoder_sequence + sequence"}
+        else:
+            common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"}
+            common_inputs["decoder_attention_mask"] = {0: "batch", 1: "decoder_sequence"}
+
+        if self.use_past:
+            self.fill_with_past_key_values_(common_inputs, direction="inputs")
+
+        return common_inputs
+
+    @property
+    def default_onnx_opset(self) -> int:
+        return 13
+


Suggested change

class T5LAOnnxConfig(OnnxSeq2SeqConfigWithPast):

@property

def inputs(self) -> Mapping[str, Mapping[int, str]]:

common_inputs = {

"input_ids": {0: "batch", 1: "encoder_sequence"},

"attention_mask": {0: "batch", 1: "encoder_sequence"},

}

if self.use_past:

common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence"

common_inputs["decoder_input_ids"] = {0: "batch"}

common_inputs["decoder_attention_mask"] = {0: "batch", 1: "past_decoder_sequence + sequence"}

else:

common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"}

common_inputs["decoder_attention_mask"] = {0: "batch", 1: "decoder_sequence"}

if self.use_past:

self.fill_with_past_key_values_(common_inputs, direction="inputs")

return common_inputs

@property

def default_onnx_opset(self) -> int:

return 13

ArthurZucker · 2025-07-15T14:06:26Z

+        )
+
+        if split_mlp_wi:
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"][


Suggested change

flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"][

flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"][

we should define a utils like a def flax_get_flatten to which you give a key and it just does this!

ArthurZucker · 2025-07-15T14:07:22Z

+        k, o, q, v = t5lax_attention_lookup(old, i, "encoder", "attention")
+        new[f"encoder.block.{i}.layer.0.layer_norm.weight"] = layer_norm
+        new[f"encoder.block.{i}.layer.0.SelfAttention.k.weight"] = k.T
+        new[f"encoder.block.{i}.layer.0.SelfAttention.o.weight"] = o.T
+        new[f"encoder.block.{i}.layer.0.SelfAttention.q.weight"] = q.T
+        new[f"encoder.block.{i}.layer.0.SelfAttention.v.weight"] = v.T
+
+        # Block i, layer 1 (MLP).
+        layer_norm = t5lax_layer_norm_lookup(old, i, "encoder", "pre_mlp_layer_norm")
+        wi, wo = t5lax_mlp_lookup(old, i, "encoder", split_mlp_wi)
+        new[f"encoder.block.{i}.layer.1.layer_norm.weight"] = layer_norm
+        if split_mlp_wi:
+            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi_0.weight"] = wi[0].T
+            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi_1.weight"] = wi[1].T
+        else:
+            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi.weight"] = wi.T
+        new[f"encoder.block.{i}.layer.1.DenseReluDense.wo.weight"] = wo.T
+
+    new["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = old[
+        "encoder/relpos_bias/rel_embedding"
+    ].T
+    new["encoder.final_layer_norm.weight"] = old["encoder/encoder_norm/scale"]
+
+    if not is_encoder_only:
+        # Decoder.
+        for i in range(num_decoder_layers):
+            # Block i, layer 0 (Self Attention).
+            layer_norm = t5lax_layer_norm_lookup(old, i, "decoder", "pre_self_attention_layer_norm")
+            k, o, q, v = t5lax_attention_lookup(old, i, "decoder", "self_attention")
+            new[f"decoder.block.{i}.layer.0.layer_norm.weight"] = layer_norm
+            new[f"decoder.block.{i}.layer.0.SelfAttention.k.weight"] = k.T
+            new[f"decoder.block.{i}.layer.0.SelfAttention.o.weight"] = o.T
+            new[f"decoder.block.{i}.layer.0.SelfAttention.q.weight"] = q.T
+            new[f"decoder.block.{i}.layer.0.SelfAttention.v.weight"] = v.T
+
+            # Block i, layer 1 (Cross Attention).
+            layer_norm = t5lax_layer_norm_lookup(old, i, "decoder", "pre_cross_attention_layer_norm")
+            k, o, q, v = t5lax_attention_lookup(old, i, "decoder", "encoder_decoder_attention")
+            new[f"decoder.block.{i}.layer.1.layer_norm.weight"] = layer_norm
+            new[f"decoder.block.{i}.layer.1.EncDecAttention.k.weight"] = k.T
+            new[f"decoder.block.{i}.layer.1.EncDecAttention.o.weight"] = o.T
+            new[f"decoder.block.{i}.layer.1.EncDecAttention.q.weight"] = q.T
+            new[f"decoder.block.{i}.layer.1.EncDecAttention.v.weight"] = v.T
+
+            # Block i, layer 2 (MLP).
+            layer_norm = t5lax_layer_norm_lookup(old, i, "decoder", "pre_mlp_layer_norm")
+            wi, wo = t5lax_mlp_lookup(old, i, "decoder", split_mlp_wi)
+            new[f"decoder.block.{i}.layer.2.layer_norm.weight"] = layer_norm
+            if split_mlp_wi:
+                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi_0.weight"] = wi[0].T
+                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi_1.weight"] = wi[1].T
+            else:
+                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi.weight"] = wi.T
+            new[f"decoder.block.{i}.layer.2.DenseReluDense.wo.weight"] = wo.T
+
+        new["decoder.final_layer_norm.weight"] = old["decoder/decoder_norm/scale"]
+        new["decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = old[
+            "decoder/relpos_bias/rel_embedding"
+        ].T
+
+        # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)


take inspiration from the convert_llama4 file for an update to this!

ArthurZucker · 2025-07-15T14:07:59Z

+
+        import numpy as np
+        import tensorflow as tf
+    except ImportError:
+        logger.error(
+            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+            "https://www.tensorflow.org/install/ for installation instructions."
+        )
+        raise
+    tf_path = os.path.abspath(tf_checkpoint_path)
+    logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+    # Load weights from TF model
+    init_vars = tf.train.list_variables(tf_path)
+    names = []
+    tf_weights = {}
+    for name, shape in init_vars:
+        logger.info(f"Loading TF weight {name} with shape {shape}")
+        array = tf.train.load_variable(tf_path, name)
+        names.append(name)
+        tf_weights[name] = array
+
+    for txt_name in names:
+        name = txt_name.split("/")
+        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
+        # which are not required for using pretrained model
+        if any(
+            n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
+            for n in name
+        ):
+            logger.info(f"Skipping {'/'.join(name)}")
+            tf_weights.pop(txt_name, None)
+            continue
+        if "_slot_" in name[-1]:
+            logger.info(f"Skipping {'/'.join(name)}")
+            tf_weights.pop(txt_name, None)
+            continue
+        pointer = model
+        array = tf_weights[txt_name]
+
+        for m_name in name:
+            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
+                scope_names = re.split(r"_(\d+)", m_name)
+            else:
+                scope_names = [m_name]
+            if scope_names[0] in ["kernel", "scale", "embedding"]:
+                pointer = getattr(pointer, "weight")
+            elif scope_names[0] == "self_attention":
+                pointer = getattr(pointer, "layer")
+                pointer = pointer[0]
+            elif scope_names[0] == "enc_dec_attention":
+                pointer = getattr(pointer, "layer")
+                pointer = pointer[1]
+            elif scope_names[0] == "dense_relu_dense":
+                pointer = getattr(pointer, "layer")
+                pointer = pointer[2]
+            elif scope_names[0] == "rms_norm":
+                if hasattr(pointer, "layer_norm"):
+                    pointer = getattr(pointer, "layer_norm")
+                elif hasattr(pointer, "final_layer_norm"):
+                    pointer = getattr(pointer, "final_layer_norm")
+            elif scope_names[0] == "scale":
+                pointer = getattr(pointer, "weight")
+            elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
+                pointer = getattr(pointer, "bias")
+            elif scope_names[0] == "squad":
+                pointer = getattr(pointer, "classifier")
+            elif scope_names[0] == "decoder" and name[1] == "logits":
+                continue
+            elif scope_names[0] == "logits":
+                pointer = getattr(pointer, "lm_head")
+            elif scope_names[0] == "wi" and len(scope_names) > 1 and scope_names[1].isdigit():
+                pointer = getattr(pointer, f"wi_{scope_names[1]}")
+                continue
+            else:
+                try:
+                    pointer = getattr(pointer, scope_names[0])
+                except AttributeError:
+                    logger.info(f"Skipping {'/'.join(name)}")
+                    continue
+            if len(scope_names) >= 2:
+                num = int(scope_names[1])
+                pointer = pointer[num]
+        if scope_names[0] not in ["kernel", "scale", "embedding"]:
+            pointer = getattr(pointer, "weight")
+        if scope_names[0] != "embedding":
+            logger.info(f"Transposing numpy weight of shape {array.shape} for {name}")
+            array = np.transpose(array)
+        try:
+            if pointer.shape != array.shape:
+                raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
+        except AssertionError as e:
+            e.args += (pointer.shape, array.shape)
+            raise
+        logger.info(f"Initialize PyTorch weight {name}")
+        pointer.data = torch.from_numpy(array.astype(np.float32))
+        tf_weights.pop(txt_name, None)
+
+    logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}.")
+    return model
+


ArthurZucker · 2025-07-15T14:08:28Z

t5 model is very old, recommending you to take inspiration from T5Gemma that was just added, and use modular to isolate the differences! 🤗

Thank you for your feedback @ArthurZucker, and also for introducing modular approach and the Gemma, I wasn't familiar with them!
I created a separate branch for my copy of gemma, available here, would you recommend that I force push it here on this branch and PR? Or should I close this and create a new PR?

You can force push here will be simpler 🫡

Sure, hopefully soon.

By the way, as Gemma's structure is relatively different from T5, when you said "t5 model is very old", did you mean in terms of architecture and design or in terms of the way it is implemented in Transformers (modular approach)? Whilst I'm happy to provide the LA version of Gemma (as is in progress in the mentioned branch), I still would need a model as close to T5 as possible, for the purpose of comparison in my paper (I guess). So what's your advice in this regard?

The best example is T5Gemma as is adheres to the recent standards of the transformers library!

The idea is just that our code in transformers is old
so has a lot of legacy

…sible

github-actions · 2025-09-19T04:31:55Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, t5la

HRezaei commented Jul 9, 2025

View reviewed changes

HRezaei force-pushed the t5la_models branch 4 times, most recently from 3aae8b0 to 06d9258 Compare July 11, 2025 16:50

HRezaei changed the title ~~[WIP] Add T5LA models~~ Add T5LA models Jul 11, 2025

ArthurZucker added the New model label Jul 15, 2025

ArthurZucker reviewed Jul 15, 2025

View reviewed changes

Add output of command add-new-model-like

1f1070a

HRezaei force-pushed the t5la_models branch 2 times, most recently from ca104a5 to 1b7338e Compare August 29, 2025 12:21

HRezaei added 4 commits September 19, 2025 05:29

Add LA model and try to fix as many unit tests and repo-checks as pos…

c643218

…sible

Remove T5LaEncoderModel

3c17431

Remove T5LaOnnxConfig

302d942

Remove backward campatibility for the new Lookahead model

c0aceaa

HRezaei force-pushed the t5la_models branch from 1b7338e to c0aceaa Compare September 19, 2025 04:30

evalstate mentioned this pull request Apr 29, 2026

Cumulative feature and defect updates from recent Transformers PRs evalstate/transformers#42

Open

	flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"][
	flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"][

Uh oh!

Conversation

HRezaei commented Jul 9, 2025

What does this PR do?

Before submitting

Who can review?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants