Code for VisionT5Model
Author: SaiMadhusudan
Pull request: no
Labels: name: Feature request

### Feature request

So right now you can't use T5 as decoder block in VisionEncoderDecoderModel , I wrote a code here which almost does that trying to get some help if it covers everything I need and can use it directly I am planning to use it for  a OCR code base 

``` python
import copy
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
from typing import Optional, Tuple, Union

from transformers import (
    PreTrainedModel,
    GenerationMixin,
    VisionEncoderDecoderConfig,
    AutoModel,
    T5Config,
    T5Stack,
    ViTModel
)
from transformers.modeling_outputs import Seq2SeqLMOutput


class VisionT5Model(PreTrainedModel, GenerationMixin):
    """
    A vision-text model using a ViT-like encoder and a T5 decoder stack.
    It mimics the design of VisionEncoderDecoderModel but replaces the decoder
    with a T5 decoder. Useful for tasks like OCR, image captioning, etc.
    """

    config_class = VisionEncoderDecoderConfig
    base_model_prefix = "vision_t5"
    main_input_name = "pixel_values"

    def __init__(self, config: VisionEncoderDecoderConfig):
        """
        Args:
            config (VisionEncoderDecoderConfig):
                Configuration for the vision-encoder–text-decoder model. 
                - config.encoder should be a vision config (e.g. ViTConfig)
                - config.decoder should be a T5Config
        """
        super().__init__(config)
        # ----------------------
        # 1) Load the Vision Encoder
        # ----------------------
        self.encoder = ViTModel(config.encoder)
        # Make sure it does NOT have a "head" for classification etc.
        if self.encoder.get_output_embeddings() is not None:
            raise ValueError("The encoder should not have a LM head; please use a bare vision backbone.")

        # ----------------------
        # 2) Build the T5 decoder stack (no encoder part!)
        # ----------------------
        # We copy the T5 config from config.decoder
        # Then ensure is_decoder=True, is_encoder_decoder=False, etc.
        t5_decoder_config = T5Config.from_dict(config.decoder.to_dict())
        t5_decoder_config.is_decoder = True
        t5_decoder_config.is_encoder_decoder = False
        t5_decoder_config.num_layers = config.decoder.num_layers
        # If you want cross-attention in T5, it must have `add_cross_attention=True`.
        # Usually T5's is_decoder implies that anyway, but just to be safe:
        t5_decoder_config.add_cross_attention = True

        self.decoder = T5Stack(t5_decoder_config)

        # Optionally, if the hidden sizes differ, we need a projection:
        if self.encoder.config.hidden_size != t5_decoder_config.d_model:
            self.enc_to_dec_proj = nn.Linear(
                self.encoder.config.hidden_size, t5_decoder_config.d_model, bias=False
            )
        else:
            self.enc_to_dec_proj = None

        # ----------------------
        # 3) Final LM head (same as T5's)
        # ----------------------
        self.lm_head = nn.Linear(t5_decoder_config.d_model, t5_decoder_config.vocab_size, bias=False)
        if t5_decoder_config.tie_word_embeddings:
            self.lm_head.weight = self.decoder.embed_tokens.weight
        self.model_dim = t5_decoder_config.d_model  # keep track if we want the T5 scaling

        # Initialize weights, etc.
        self.post_init()

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def get_input_embeddings(self):
        """By convention, the 'input embeddings' come from the decoder if needed."""
        return self.decoder.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.decoder.set_input_embeddings(new_embeddings)

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(
        self,
        pixel_values: torch.FloatTensor,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.BoolTensor] = None,
        encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **decoder_kwargs
    ) -> Union[Seq2SeqLMOutput, Tuple[torch.FloatTensor]]:
        """
        pixel_values: (batch, channels, height, width)
            The images to encode (e.g. from ViTFeatureExtractor).
        decoder_input_ids: (batch, tgt_seq_len)
            Input tokens to the T5 decoder.
        labels: (batch, tgt_seq_len)
            If given, we compute LM loss by teacher-forcing and produce CrossEntropyLoss.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        use_cache = use_cache if use_cache is not None else self.config.decoder.use_cache

        # 1) Run the vision encoder if needed
        if encoder_outputs is None:
            encoder_outputs = self.encoder(pixel_values=pixel_values, return_dict=True)
        # encoder_outputs.last_hidden_state shape => (batch, seq_len, hidden_size)
        hidden_states = encoder_outputs.last_hidden_state

        # Possibly project to match T5 dimension
        if self.enc_to_dec_proj is not None:
            hidden_states = self.enc_to_dec_proj(hidden_states)

        # 2) Prepare decoder inputs
        # If we have labels but no decoder_input_ids, shift-right internally
        if labels is not None and decoder_input_ids is None:
            # Standard T5 shift-right:
            decoder_input_ids = self._shift_right(labels)

        # T5 decoder forward
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=hidden_states,
            encoder_attention_mask=None,  # If you want to mask out padding in hidden_states, pass it here
            past_key_values=past_key_values,
            use_cache=use_cache,
            return_dict=True,
            **decoder_kwargs,
        )

        sequence_output = decoder_outputs[0]  # (batch, tgt_len, d_model)
        # 3) Final LM head
        # T5 typically scales by d_model^-0.5 if tie_word_embeddings = True, 
        # but you can do that if needed. 
        if self.config.decoder.tie_word_embeddings:
            sequence_output = sequence_output * (self.model_dim ** -0.5)
        logits = self.lm_head(sequence_output)

        loss = None
        if labels is not None:
            # Compute standard LM loss
            loss_fct = CrossEntropyLoss(ignore_index=-100)
            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))

        if not return_dict:
            # Return (loss, logits, past, decoder_outputs, encoder_outputs)
            out = (logits,) + decoder_outputs[1:] + (encoder_outputs,)
            return ((loss,) + out) if loss is not None else out

        return Seq2SeqLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=hidden_states,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        past_key_values=None,
        encoder_outputs=None,
        **kwargs,
    ):
        """
        During generation, the `generate()` method calls this to assemble the inputs for each step.
        """
        if past_key_values is not None:
            # we only need to pass the last token of decoder_input_ids
            decoder_input_ids = decoder_input_ids[:, -1:].clone()

        return {
            "pixel_values": None,  # not needed if `encoder_outputs` is already computed
            "decoder_input_ids": decoder_input_ids,
            "past_key_values": past_key_values,
            "encoder_outputs": encoder_outputs,
            "use_cache": kwargs.get("use_cache"),
        }
        
    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor) -> torch.Tensor:
        return self._shift_right(labels)

    def _reorder_cache(self, past_key_values, beam_idx):
        # if decoder past is not included in output
        # speedy decoding is disabled and no need to reorder
        if past_key_values is None:
            print("You might want to consider setting `use_cache=True` to speed up decoding")
            return past_key_values

        reordered_decoder_past = ()
        for layer_past_states in past_key_values:
            # get the correct batch idx from layer past batch dim
            # batch dim of `past` is at 2nd position
            reordered_layer_past_states = ()
            for layer_past_state in layer_past_states:
                # need to set correct `past` for each of the four key / value states
                reordered_layer_past_states = reordered_layer_past_states + (
                    layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)),
                )

            if reordered_layer_past_states[0].shape != layer_past_states[0].shape:
                raise ValueError(
                    f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched"
                )
            if len(reordered_layer_past_states) != len(layer_past_states):
                raise ValueError(
                    f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched"
                )

            reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,)
        return reordered_decoder_past


        
    def _shift_right(self, labels: torch.LongTensor) -> torch.LongTensor:
        """
        Same shifting that T5 does: pad -> start token -> ... -> y[0..-2]
        """
        # In T5, the decoder_start_token_id is often the same as pad_token_id
        # But check or override as needed.
        decoder_start_token_id = self.config.decoder.decoder_start_token_id
        if decoder_start_token_id is None:
            # default fallback
            decoder_start_token_id = self.config.decoder.pad_token_id

        pad_token_id = self.config.decoder.pad_token_id
        # create shifted ids
        shifted = labels.new_zeros(labels.shape)
        shifted[..., 1:] = labels[..., :-1].clone()
        shifted[..., 0] = decoder_start_token_id
        # replace -100 with pad_token_id
        shifted.masked_fill_(shifted == -100, pad_token_id)
        return shifted
```

### Motivation

For OCR Project

### Your contribution

T5 can be used a decoder block for vision models 

Comments:
Hi @SaiMadhusudan , these questions are usually more appropriate for the Discord or forums than Github issues, since it's not an actual bug! http://hf.co/join/discord