uform-gen2-dpo / modeling_uform_gen.py
VoVoR's picture
Add: DPO model card init
0556e7f
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")