|
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: |
|
|
|
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 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") |
|
|
