from typing import List, Optional, Tuple, Union
import warnings, os, torch
import torch.nn as nn
from transformers import AutoConfig, PretrainedConfig, PreTrainedModel, AutoModelForCausalLM, AutoTokenizer
from transformers.modeling_utils import ContextManagers, no_init_weights
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation.utils import GenerateOutput
from .configuration_apollo import ApolloConfig
from .vision_tower import ApolloVisionTower
from .mm_connector import MMConnector
IGNORE_INDEX = -100
X_TOKEN_INDEX = -200
def get_model_config(config):
default_keys = ["llm_cfg", "vision_tower_cfg", "mm_connector_cfg"]
if hasattr(config, "_name_or_path") and len(config._name_or_path) >= 2:
root_path = config._name_or_path
else:
root_path = config.resume_path
return_pths = []
for key in default_keys:
cfg = getattr(config, key, None)
if isinstance(cfg, dict):
try:
return_pths.append(os.path.join(root_path, key[:-4]))
except:
raise ValueError(f"Cannot find resume path in config for {key}!")
elif isinstance(cfg, PretrainedConfig):
return_pths.append(os.path.join(root_path, key[:-4]))
elif isinstance(cfg, str):
return_pths.append(cfg)
return_list = []
for pth in return_pths:
return_list.append(AutoConfig.from_pretrained(pth, trust_remote_code=True))
return return_list
def build_llm_and_tokenizer(
llm_cfg: str,
config: PretrainedConfig,
attn_implementation=None,
model_max_length=None,
*args,
**kwargs,
) -> PreTrainedModel:
llm_arch = getattr(llm_cfg, "architectures")[0].lower()
llm_path = llm_cfg._name_or_path
llm = AutoModelForCausalLM.from_pretrained(
llm_path, config=llm_cfg, torch_dtype=eval(config.model_dtype), *args, **kwargs
)
tokenizer = AutoTokenizer.from_pretrained(
llm_path,
model_max_length=llm_cfg.model_max_length,
padding_side="right",
use_fast=False,
legacy=False,
**kwargs
)
#config.hidden_size = llm.config.hidden_size
return llm, tokenizer
class ApolloForCausalLM(PreTrainedModel):
def __init__(self, config: ApolloConfig, *args, **kwargs):
super().__init__(config)
llm_cfg, vision_tower_cfg, mm_connector_cfg = get_model_config(config)
model_dtype = getattr(config, "model_dtype", "torch.float16")
if not hasattr(config, "model_dtype"):
warnings.warn("model_dtype not found in config, defaulting to torch.float16.")
config.model_dtype = model_dtype
# Initialize weights and apply final processing
self.lm_head = nn.Linear(llm_cfg.hidden_size, config.vocab_size, bias=False)
self.vision_tower = ApolloVisionTower(config, vision_tower_cfg)
self.mm_connector = MMConnector.from_pretrained(mm_connector_cfg._name_or_path)
self.llm, self.tokenizer = build_llm_and_tokenizer(llm_cfg, config, *args, **kwargs)
self.post_init()
self.is_loaded = True
def forward(
self,
input_ids: torch.LongTensor = 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,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
vision_input: Optional[List[torch.FloatTensor]] = None,
data_types: Optional[List[str]] = None,
return_dict: Optional[bool] = None,
cache_position=None,
) -> Union[Tuple, CausalLMOutputWithPast]:
if inputs_embeds is None:
(
input_ids,
position_ids,
attention_mask,
past_key_values,
inputs_embeds,
labels
) = self.prepare_inputs_labels_for_multimodal(
input_ids,
position_ids,
attention_mask,
past_key_values,
labels,
vision_input,
data_types
)
return self.get_llm().forward(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
labels=labels,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
@torch.no_grad()
def generate(
self,
inputs: Optional[torch.Tensor] = None,
vision_input: Optional[List[torch.Tensor]] = None,
data_types: Optional[List[str]] = None,
**kwargs,
) -> Union[GenerateOutput, torch.LongTensor]:
position_ids = kwargs.pop("position_ids", None)
attention_mask = kwargs.pop("attention_mask", None)
if "inputs_embeds" in kwargs:
raise NotImplementedError("`inputs_embeds` is not supported")
if vision_input is not None:
(inputs, position_ids, attention_mask, _, inputs_embeds, _) = self.prepare_inputs_labels_for_multimodal(
inputs, position_ids, attention_mask, None, None, vision_input, data_types=data_types)
else:
inputs_embeds = self.embed_tokens(inputs)
return self.get_llm().generate(position_ids=position_ids, attention_mask=attention_mask,
inputs_embeds=inputs_embeds, **kwargs)
def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
vision_input = kwargs.pop("vision_input", None)
data_types = kwargs.pop("data_types", None)
inputs = self.get_llm().prepare_inputs_for_generation(input_ids, past_key_values=past_key_values,
inputs_embeds=inputs_embeds, **kwargs)
if vision_input is not None:
inputs["vision_input"] = vision_input
if data_types is not None:
inputs["data_types"] = data_types
return inputs
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
*model_args,
config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
cache_dir: Optional[Union[str, os.PathLike]] = None,
ignore_mismatched_sizes: bool = False,
force_download: bool = False,
local_files_only: bool = False,
token: Optional[Union[str, bool]] = None,
revision: str = "main",
use_safetensors: bool = None,
**kwargs,
):
return cls.load_pretrained(
pretrained_model_name_or_path,
*model_args,
config=config,
cache_dir=cache_dir,
ignore_mismatched_sizes=ignore_mismatched_sizes,
force_download=force_download,
local_files_only=local_files_only,
token=token,
revision=revision,
use_safetensors=use_safetensors,
**kwargs,
)
def get_llm(self):
return self.llm
def get_vision_tower(self):
return self.vision_tower
def get_mm_connector(self):
return self.mm_connector
@classmethod
def load_pretrained(cls, model_path_or_config, *args, **kwargs):
kwargs.pop("config", None)
if isinstance(model_path_or_config, str):
config = AutoConfig.from_pretrained(model_path_or_config, trust_remote_code=True, **kwargs)
elif isinstance(model_path_or_config, ApolloConfig):
config = model_path_or_config
else:
raise NotImplementedError(f"wrong type, {type(model_path_or_config)} \
{isinstance(model_path_or_config, ApolloConfig)}")
model_dtype = getattr(config, "model_dtype", "torch.float16")
if not hasattr(config, "model_dtype"):
warnings.warn("model_dtype not found in config, defaulting to torch.float16.")
config.model_dtype = model_dtype
with ContextManagers([no_init_weights(_enable=True), ]):
vlm = cls(config, *args, **kwargs)
if hasattr(vlm, "llm") and hasattr(vlm, "vision_tower") and hasattr(vlm, "mm_connector"):
if vlm.is_loaded:
return vlm
else:
print('loading model failed!')
else:
print('loading model failed!')
def _encode_mm(self, x):
x = self.get_vision_tower()(x)
x = self.mm_connector(x)
return x
def encode_mm_minibatch(self, x):
split_sizes = [x_s[0].shape[0] for x_s in x]
x = [torch.split(torch.cat([x_s[i] for x_s in x], dim=0), self.config.encode_batch_size) for i in
range(self.get_vision_tower().num_vision_encoders)]
swapped_x = []
for i in range(len(x[0])):
swapped_x.append([x_s[i] for x_s in x])
features = []
for xx in swapped_x:
xx = self._encode_mm(xx)
features.append(xx)
x = torch.cat(features, dim=0)
x = torch.split(x, split_sizes, dim=0)
return [xx.contiguous().view(-1, xx.shape[2]) for xx in x]
def prepare_inputs_labels_for_multimodal(
self, input_ids, position_ids, attention_mask, past_key_values, labels, vision_input, data_types
):
vision_tower = self.get_vision_tower()
if vision_tower is None or vision_input is None or input_ids.shape[1] == 1:
if (
past_key_values is not None
and vision_tower is not None
and vision_input is not None
and input_ids.shape[1] == 1
):
target_shape = past_key_values[-1][-1].shape[-2] + 1
attention_mask = torch.cat(
(
attention_mask,
torch.ones(
(
attention_mask.shape[0],
target_shape - attention_mask.shape[1],
),
dtype=attention_mask.dtype,
device=attention_mask.device,
),
),
dim=1,
)
position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
return (
input_ids,
position_ids,
attention_mask,
past_key_values,
None,
labels,
)
'''
vision_input is a list of tuples, and data_type is a list of strings:
data_type = ['image', 'video', 'video'..., 'text']
(for one video and two image encoders)
vision_input =
[
[image(1, T, C, H, W), image(1, T, C, H, W), image(1, T, C, H, W)],
[video(Nc1, C, T, H, W), video(Nc1, T, C, H, W), video(Nc1, T, C, H, W)],
[video(Nc2, C, T, H, W), video(Nc2, T, C, H, W), video(Nc2, T, C, H, W)],
]
-> video encoders typlically expect (C,T,H,W), images expect (C,H,W).
'''
# ====================================================================================================
merged_mm_features = self.encode_mm_minibatch(vision_input)
if not getattr(self.config, "tune_language_model", True) and getattr(self.config, "use_mm_start_end", False):
raise NotImplementedError
# ====================================================================================================
# Let's just add dummy tensors if they do not exist,
# it is a headache to deal with None all the time.
# But it is not ideal, and if you have a better idea,
# please open an issue / submit a PR, thanks.
_labels = labels
_position_ids = position_ids
_attention_mask = attention_mask
if attention_mask is None:
attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
else:
attention_mask = attention_mask.bool()
if position_ids is None:
position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
if labels is None:
labels = torch.full_like(input_ids, IGNORE_INDEX)
# remove the padding using attention_mask
input_ids_copy = input_ids.clone()
# kentang-mit@: Otherwise tokenizer out of bounds. Embeddings of image tokens will not be used.
input_ids_copy[input_ids_copy == X_TOKEN_INDEX] = 0
input_embeds = self.get_llm().model.embed_tokens(input_ids_copy)
input_ids = [
cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)
]
input_embeds_1 = [
cur_input_embeds[cur_attention_mask]
for cur_input_embeds, cur_attention_mask in zip(input_embeds, attention_mask)
]
labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]
# input_ids, new_input_embeds = self.inputs_merger(input_ids, input_embeds_1, merged_mm_features)
new_labels = []
new_input_embeds = []
# print("BEFORE BATCH LOOP:", len(input_ids), input_ids[0].shape, input_ids[0].device, [(x == X_TOKEN_INDEX).sum() for x in input_ids])
# kentang-mit@: If some part of the model is executed in the loop, the the loop length needs to be a constant.
for batch_idx, (cur_labels, cur_input_ids, mm_features) in enumerate(
zip(labels, input_ids, merged_mm_features)):
cur_input_ids = input_ids[batch_idx]
num_mm = (cur_input_ids == X_TOKEN_INDEX).sum()
if num_mm == 0:
cur_input_embeds_1 = input_embeds_1[batch_idx]
cur_input_embeds = torch.cat([cur_input_embeds_1, mm_features[0:0]], dim=0)
new_input_embeds.append(cur_input_embeds)
new_labels.append(cur_labels)
# kenang-mit@: we do not have placeholdr image for text-only data now.
continue
if mm_features.shape[0] != num_mm:
print(data_types[batch_idx])
assert num_mm == len(
mm_features), f'Error in {data_types[batch_idx]}{num_mm}=/={len(mm_features)} not the same number of vision tokens in and vision embeddings!'
cur_input_embeds = input_embeds_1[batch_idx]
image_token_indices = (
[-1] + torch.where(cur_input_ids == X_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
)
cur_input_ids_noim = []
cur_labels = labels[batch_idx]
cur_labels_noim = []
cur_input_embeds_no_im = []
for i in range(len(image_token_indices) - 1):
cur_input_ids_noim.append(cur_input_ids[image_token_indices[i] + 1: image_token_indices[i + 1]])
cur_labels_noim.append(cur_labels[image_token_indices[i] + 1: image_token_indices[i + 1]])
cur_input_embeds_no_im.append(cur_input_embeds[image_token_indices[i] + 1: image_token_indices[i + 1]])
cur_new_input_embeds = []
cur_new_labels = []
for i in range(num_mm + 1):
cur_new_input_embeds.append(cur_input_embeds_no_im[i])
# print("cur_new_input_embeds1", cur_new_input_embeds.shape[-1])
cur_new_labels.append(cur_labels_noim[i])
if i < num_mm:
cur_image_features = mm_features[i:i + 1]
cur_new_input_embeds.append(cur_image_features)
# print("cur_new_input_embeds2", cur_new_input_embeds.shape[-1])
cur_new_labels.append(
torch.full(
(cur_image_features.shape[0],),
IGNORE_INDEX,
device=cur_labels.device,
dtype=cur_labels.dtype,
)
)
cur_new_input_embeds = torch.cat(cur_new_input_embeds)
cur_new_labels = torch.cat(cur_new_labels)
new_input_embeds.append(cur_new_input_embeds)
new_labels.append(cur_new_labels)
# Truncate sequences to max length as image embeddings can make the sequence longer
tokenizer_model_max_length = getattr(self.get_llm().config, "tokenizer_model_max_length", None)
if tokenizer_model_max_length is not None:
if any(len(x) > tokenizer_model_max_length for x in new_input_embeds):
priny("Inputs truncated!")
new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
# Combine them
max_len = max(x.shape[0] for x in new_input_embeds)
batch_size = len(new_input_embeds)
new_input_embeds_padded = []
new_labels_padded = torch.full(
(batch_size, max_len),
IGNORE_INDEX,
dtype=new_labels[0].dtype,
device=new_labels[0].device,
)
attention_mask = torch.zeros(
(batch_size, max_len),
dtype=attention_mask.dtype,
device=attention_mask.device,
)
position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)
for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):
cur_len = cur_new_embed.shape[0]
if getattr(self.get_llm().config, "tokenizer_padding_side", "right") == "left":
new_input_embeds_padded.append(
torch.cat(
(
torch.zeros(
(max_len - cur_len, cur_new_embed.shape[1]),
dtype=cur_new_embed.dtype,
device=cur_new_embed.device,
),
cur_new_embed,
),
dim=0,
)
)
if cur_len > 0:
new_labels_padded[i, -cur_len:] = cur_new_labels
attention_mask[i, -cur_len:] = True
position_ids[i, -cur_len:] = torch.arange(
0, cur_len, dtype=position_ids.dtype, device=position_ids.device
)
else:
new_input_embeds_padded.append(
torch.cat(
(
cur_new_embed,
torch.zeros(
(max_len - cur_len, cur_new_embed.shape[1]),
dtype=cur_new_embed.dtype,
device=cur_new_embed.device,
),
),
dim=0,
)
)
if cur_len > 0:
new_labels_padded[i, :cur_len] = cur_new_labels
attention_mask[i, :cur_len] = True
position_ids[i, :cur_len] = torch.arange(
0, cur_len, dtype=position_ids.dtype, device=position_ids.device
)
new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
if _labels is None:
new_labels = None
else:
new_labels = new_labels_padded
if _attention_mask is None:
attention_mask = None
else:
attention_mask = attention_mask.to(dtype=_attention_mask.dtype)
if _position_ids is None:
position_ids = None
return (
None,
position_ids,
attention_mask,
past_key_values,
new_input_embeds,
new_labels,
)