modeling_uform_gen.py

from typing import List, Optional, Tuple, Union

from .configuration_uform_gen import VLMConfig

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from torch import nn

from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.models.auto.modeling_auto import AutoModelForCausalLM, AutoModel
from transformers import AutoConfig
from transformers.utils import logging

from .vision_encoder import VisionEncoder


class ImageFeaturesPooler(nn.Module):
    def __init__(self, config, text_config):
        super().__init__()
        self.pooler = nn.TransformerDecoderLayer(
            config.image_encoder_hidden_size,
            config.image_pooler_num_attn_heads,
            config.image_pooler_intermediate_size,
            activation=nn.functional.silu,
            batch_first=True,
            norm_first=True,
        )
        self.image_latents = nn.Parameter(
            torch.randn(1, config.num_image_latents, config.image_encoder_hidden_size)
            * config.initializer_range**0.5
        )
        self.projection = nn.Linear(config.image_encoder_hidden_size, text_config.hidden_size)

    def forward(self, features):
        features = self.pooler(
            self.image_latents.expand(features.size(0), -1, -1), features
        )

        return self.projection(features)


class VLMPreTrainedModel(PreTrainedModel):
    config_class = VLMConfig
    base_model_prefix = "vlm"
    supports_gradient_checkpointing = True
    _no_split_modules = []
    _skip_keys_device_placement = "past_key_values"

    def _init_weights(self, module):
        pass

    def _initialize_weights(self, module):
        pass


class VLMForCausalLM(VLMPreTrainedModel):
    def __init__(self, config: VLMConfig):
        super().__init__(config)

        self.config = config
        self.text_config = AutoConfig.from_pretrained(
            config.text_decoder_name_or_path,
            trust_remote_code=True
        )
        
        self.text_decoder = AutoModelForCausalLM.from_config(
            self.text_config,
            trust_remote_code=True
        )

        self.image_encoder = VisionEncoder(
            config.image_encoder_hidden_size,
            config.image_encoder_patch_size,
            config.image_encoder_num_layers,
            config.image_encoder_num_heads,
        )

        self.image_pooler = ImageFeaturesPooler(config, self.text_config)

    def get_input_embeddings(self):
        return self.text_decoder.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.text_decoder.set_input_embeddings(value)

    def get_images_embeddings(self, images):
        features = self.image_encoder(images)
        return self.image_pooler(features)

    def gather_continuous_embeddings(
        self,
        input_ids: torch.Tensor,
        word_embeddings: torch.Tensor,
        image_embeddings: torch.Tensor
    ) -> torch.Tensor:
        
        start_indices = (input_ids == self.config.image_token_id).nonzero()[:, 1]
        embeddings = []
        for sample_idx, start_idx in enumerate(start_indices.tolist()):
            embeddings.append(
                torch.cat(
                    (
                        word_embeddings[sample_idx, :start_idx],
                        image_embeddings[sample_idx],
                        word_embeddings[sample_idx, start_idx + 1 :],
                    ),
                    dim=0,
                )
            )

        return torch.stack(embeddings, dim=0)

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        images: torch.Tensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None
    ) -> Union[dict, Tuple, CausalLMOutputWithPast]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time"
            )
        elif input_ids is None and inputs_embeds is None:
            raise ValueError("You have to specify either input_is or inputs_embeds")

        if inputs_embeds is None and past_key_values is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

            if images is not None:
                image_embeds = self.get_images_embeddings(images)
                inputs_embeds = self.gather_continuous_embeddings(
                    input_ids,
                    inputs_embeds,
                    image_embeds
                )

        if position_ids is None:
            seq_length = (
                inputs_embeds.shape[1]
                if inputs_embeds is not None
                else input_ids.shape[1]
            )
            past_key_values_length = 0

            if past_key_values is not None:
                past_key_values_length = past_key_values[0][0].shape[2]

            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length,
                seq_length + past_key_values_length,
                dtype=torch.long,
                device=device,
            )
            position_ids = position_ids.unsqueeze(0)

        outputs = self.text_decoder(
            inputs_embeds=inputs_embeds,
            input_ids=input_ids if past_key_values is not None else None,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            use_cache=use_cache,
            return_dict=return_dict,
        )

        return outputs

    def prepare_inputs_for_generation(
        self,
        input_ids,
        images=None,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]

        position_ids = kwargs.get("position_ids", None)
        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
            n_samples = inputs_embeds.shape[0]
        else:
            model_inputs = {"input_ids": input_ids}
            n_samples = input_ids.shape[0]

        if images is not None:
            model_inputs["images"] = images

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "images": images if past_key_values is None else None,
            }
        )
        return model_inputs

    @classmethod
    def from_config(cls, config, **kwargs):
        return cls._from_config(config, **kwargs)


VLMConfig.register_for_auto_class()
VLMForCausalLM.register_for_auto_class("AutoModel")