from collections import Counter, defaultdict
import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor
from torch import nn
from torch.nn import CrossEntropyLoss
import copy
import math
from transformers.activations import gelu
from typing import List, Optional, Tuple, Union
from transformers.modeling_utils import PreTrainedModel, PretrainedConfig
from transformers import CONFIG_MAPPING
from transformers.modeling_outputs import BaseModelOutput
from transformers import GenerationConfig
from transformers import CLIPConfig, CLIPProcessor, CLIPModel, AutoModel
from transformers import WhisperConfig, WhisperPreTrainedModel, WhisperModel
from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig
def most_frequent_element(tensor):
flattened_list = tensor.flatten().tolist()
counter = Counter(flattened_list)
most_common_element = counter.most_common(1)[0][1]
return most_common_element
class MM_LLMs_Config(PretrainedConfig):
model_type = 'mm_llms'
is_composition = True
def __init__(
self,
image_config=None,
llm_config=None,
vision_select_layer=None,
**kwargs
):
self.image_config = image_config
self.llm_config = llm_config
self.vision_select_layer = vision_select_layer
if isinstance(self.image_config, dict):
image_config["model_type"] = (
image_config["model_type"] if "model_type" in image_config else "clip"
)
self.image_config = CONFIG_MAPPING[image_config["model_type"]](**image_config)
if isinstance(self.llm_config, dict):
llm_config["model_type"] = llm_config["model_type"] if "model_type" in llm_config else "llama"
self.llm_config = CONFIG_MAPPING[llm_config["model_type"]](**llm_config)
super().__init__(**kwargs)
class LlavaMultiModalProjector(nn.Module):
def __init__(self, in_hidden_size, out_hidden_size, conv_kernel=None, conv_stride=3):
super().__init__()
self.conv_kernel = conv_kernel
if conv_kernel:
self.linear_1 = nn.Conv1d(
in_hidden_size,
out_hidden_size,
kernel_size=conv_kernel,
stride=conv_stride)
else:
self.linear_1 = nn.Linear(
in_hidden_size,
out_hidden_size,
bias=True,
)
self.act = gelu
self.linear_2 = nn.Linear(
out_hidden_size,
out_hidden_size,
bias=True
)
def forward(self, image_features):
hidden_states = self.linear_1(image_features)
if self.conv_kernel:
hidden_states = hidden_states.transpose(1, 2).contiguous()
hidden_states = self.act(hidden_states)
hidden_states = self.linear_2(hidden_states)
return hidden_states
class MM_LLMs(PreTrainedModel):
config_class = MM_LLMs_Config
supports_gradient_checkpointing = True
_supports_flash_attn_2 = True
def __init__(self, config, flash_attention=False, dtype=torch.float32):
super().__init__(config)
self.config = config
self.image_encoder = AutoModel.from_config(config.image_config)
self.llm = AutoModelForCausalLM.from_config(
config.llm_config,
use_flash_attention_2=flash_attention,
torch_dtype=dtype,
)
self.image_projector = LlavaMultiModalProjector(
config.image_config.vision_config.hidden_size,
config.llm_config.hidden_size
)
def forward(self,
input_ids: torch.LongTensor = None,
image_index: torch.LongTensor = None,
audio_index: torch.LongTensor = None,
image_starts: torch.int = None,
image_ends: torch.int = None,
audio_starts: torch.int = None,
audio_ends: torch.int = None,
images: torch.FloatTensor = None,
audios: torch.FloatTensor = 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,
labels: Optional[torch.LongTensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
use_cache: Optional[bool] = None,
return_dict: Optional[bool] = None, **kwargs):
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
images = images.type(self.image_encoder.dtype) if images is not None else None
audios = audios.type(self.audio_encoder.dtype) if audios is not None else None
model_inputs = self.prepare_inputs_for_generation(
input_ids=input_ids,
image_index=image_index,
audio_index=audio_index,
image_starts=image_starts,
image_ends=image_ends,
audio_starts=audio_starts,
audio_ends=audio_ends,
images=images,
audios=audios,
attention_mask=attention_mask,
labels=labels)
outputs = self.llm(
inputs_embeds=model_inputs['inputs_embeds'],
attention_mask=model_inputs['attention_mask'],
labels=model_inputs['labels'],
return_dict=return_dict)
return outputs
def prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
inputs_embeds=None,
images=None,
audios=None,
audio_starts=None,
audio_ends=None,
image_starts=None,
image_ends=None,
attention_mask=None,
labels=None,
audio_index=None,
image_index=None,
**kwargs):
image_features = self.encode_image(
images) if images is not None else None
embed_tokens = self.llm.model.embed_tokens
text_embeddings = embed_tokens(input_ids)
batch_size = text_embeddings.shape[0]
seq_len = text_embeddings.shape[1]
embed_dim = text_embeddings.shape[2]
max_count_image = most_frequent_element(image_index)
seq_image = image_features.shape[1]
new_len = text_embeddings.shape[1] + seq_image * max_count_image
final_embedding = torch.zeros(
batch_size, new_len, embed_dim,
device=text_embeddings.device,
dtype=text_embeddings.dtype
)
final_embedding[:, :seq_len] = text_embeddings
final_attention_mask = torch.zeros(
batch_size, new_len,
device=attention_mask.device,
dtype=attention_mask.dtype
)
final_attention_mask[:, :seq_len] = attention_mask
if labels is not None:
final_labels = torch.full(
(batch_size, new_len),
-100,
device=labels.device,
dtype=labels.dtype
)
final_labels[:, :seq_len] = labels
else:
final_labels = None
image_id = int(image_starts[0])
where_is = torch.where(input_ids == image_id)
positions = defaultdict(int)
b_image = 0
for i in range(len(where_is[0])):
b, k = where_is[0][i], where_is[1][i]
int_b = int(b)
int_k = int(k)
l = int(input_ids[b, k])
f = image_features[b_image]
b_image += 1
c = torch.cat([final_embedding[b, :int_k + 1 + positions[int_b]],
f, text_embeddings[b, k + 1:]])
final_embedding[b, :len(c)] = c
final_attention_mask[b, :len(c)] = 1.0
if labels is not None:
ignore = torch.tensor([-100] * len(f), device=labels.device)
c_label = torch.cat(
[final_labels[b, :int_k + 1 + positions[int_b]], ignore, labels[b, k + 1:]])
final_labels[b, :len(c)] = c_label
positions[int_b] += len(f)
model_inputs = {
"input_ids": input_ids,
"inputs_embeds": final_embedding,
"use_cache": kwargs.get("use_cache"),
"attention_mask": final_attention_mask,
"labels": final_labels,
}
return model_inputs
def encode_image(self, images):
if self.config.vision_select_layer is not None:
encoded = self.image_encoder.vision_model(images, output_hidden_states=True)
encoded = encoded.hidden_states[self.config.vision_select_layer]
else:
encoded = self.image_encoder.vision_model(images)[0]
image_features = self.image_projector(encoded)
return image_features