import copy
import os
from typing import Callable, List, Optional, Tuple, Union
import numpy as np
import torch
from torch.nn import CrossEntropyLoss
from torch.nn.utils.rnn import pad_sequence
import warnings
from torch import Tensor, nn
from transformers import (
PreTrainedModel,
PreTrainedTokenizer,
Blip2VisionModel,
Blip2QFormerModel,
Blip2Model,
Blip2PreTrainedModel,
Blip2ForConditionalGeneration,
GenerationConfig,
)
from transformers.models.blip_2.modeling_blip_2 import (
Blip2ForConditionalGenerationModelOutput,
)
from transformers.utils import logging
from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList
from .modeling_chatglm import (
ChatGLMForConditionalGeneration,
InvalidScoreLogitsProcessor,
)
from .configuration_blip2chatglm import Blip2ChatGLMConfig
logger = logging.get_logger(__name__)
class Blip2ChatGLMForConditionalGeneration(Blip2ForConditionalGeneration):
config_class = Blip2ChatGLMConfig
def __init__(self, config: Blip2ChatGLMConfig):
Blip2PreTrainedModel.__init__(self, config)
# NOTE: we only initialize Blip2PreTrainedModel
# directly call super().__init__() will cause error since ChatGLM cannot be found by AutoModel
self.vision_model = Blip2VisionModel(config.vision_config)
self.query_tokens = nn.Parameter(
torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size)
)
self.qformer = Blip2QFormerModel(config.qformer_config)
self.language_projection = nn.Linear(
config.qformer_config.hidden_size, config.text_config.hidden_size
)
self.language_model = ChatGLMForConditionalGeneration(config.text_config)
# Initialize weights and apply final processing
# self.post_init()
def setup_dtype(self, vision_encoder_dtype: str = "fp32", lm_dtype: str = "fp16"):
if vision_encoder_dtype == "fp32":
self.vision_model = self.vision_model.float()
elif vision_encoder_dtype == "fp16":
self.vision_model = self.vision_model.half()
else:
raise NotImplementedError(
f"Unsupported vision_encoder_dtype: {vision_encoder_dtype}"
)
if lm_dtype == "fp32":
self.language_model = self.language_model.float()
elif lm_dtype == "fp16":
self.language_model = self.language_model.half()
elif lm_dtype == "int4":
self.language_model = self.language_model.half().quantize(4)
elif lm_dtype == "int8":
self.language_model = self.language_model.half().quantize(8)
else:
raise NotImplementedError(f"Unsupported lm_dtype: {lm_dtype}")
def forward(
self,
pixel_values: torch.FloatTensor,
input_ids: torch.FloatTensor,
image_slot_offset: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.LongTensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
labels: Optional[torch.LongTensor] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple, Blip2ForConditionalGenerationModelOutput]:
"""_summary_
Args:
pixel_values (torch.FloatTensor): _description_
input_ids (torch.FloatTensor): input_ids[:, :num_query_tokens] should be filled with tokenizer.unk_token_id
image_slot_offset (Optional[torch.LongTensor], optional): if not set, all vtokens are placed as prefix (image_slot_offset = torch.zeros(bsz)). Defaults to None.
attention_mask (Optional[torch.LongTensor], optional): _description_. Defaults to None.
output_attentions (Optional[bool], optional): _description_. Defaults to None.
output_hidden_states (Optional[bool], optional): _description_. Defaults to None.
labels (Optional[torch.LongTensor], optional): _description_. Defaults to None.
return_dict (Optional[bool], optional): _description_. Defaults to None.
Returns:
Union[Tuple, Blip2ForConditionalGenerationModelOutput]: _description_
"""
return_dict = (
return_dict if return_dict is not None else self.config.use_return_dict
)
# step 1: forward the images through the vision encoder,
# to get image embeddings of shape (batch_size, seq_len, hidden_size)
vision_outputs = self.vision_model(
pixel_values=pixel_values,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
image_embeds = vision_outputs[0]
# step 2: forward the query tokens through the QFormer, using the image embeddings for cross-attention
image_attention_mask = torch.ones(
image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device
)
query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
query_outputs = self.qformer(
query_embeds=query_tokens,
encoder_hidden_states=image_embeds,
encoder_attention_mask=image_attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
query_output = query_outputs[0]
# step 3: use the language model, conditioned on the query outputs and the prompt
language_model_inputs = self.language_projection(query_output)
inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
if image_slot_offset is None:
# image as prefix
# update data to avoid inplace operation of leaf Variable
inputs_embeds.data[
:, : self.config.num_query_tokens, :
] = language_model_inputs
else:
for i, offset in enumerate(image_slot_offset):
inputs_embeds.data[
i, offset : offset + self.config.num_query_tokens, :
] = language_model_inputs[i]
outputs = self.language_model(
input_ids=input_ids,
inputs_embeds=inputs_embeds,
attention_mask=attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
logits = outputs.logits if return_dict else outputs[0]
loss = None
# we compute the loss here since we need to take into account the sequence length of the query embeds
if labels is not None:
logits = logits[:, -labels.size(1) :, :]
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous().to(logits.device)
# Flatten the tokens
loss_fct = CrossEntropyLoss(reduction="mean")
loss = loss_fct(
shift_logits.view(-1, self.config.text_config.vocab_size),
shift_labels.view(-1),
)
if not return_dict:
output = (logits, vision_outputs, query_outputs, outputs)
return ((loss,) + output) if loss is not None else output
return Blip2ForConditionalGenerationModelOutput(
loss=loss,
logits=logits,
vision_outputs=vision_outputs,
qformer_outputs=query_outputs,
language_model_outputs=outputs,
)
def prepare_inputs_for_chat(
self,
tokenizer: PreTrainedTokenizer,
batch_messages: List[List[Tuple[str, str, List[Tuple[torch.Tensor, int]]]]],
max_length: int,
user_role: str = "问",
bot_role: str = "答",
):
device = self.device
nvtokens = self.config.num_query_tokens
# 1. Prepare token ids
all_images = []
all_image_slots = []
all_input_ids = []
for messages in batch_messages:
images = []
image_slots = []
input_ids = []
round_roles = [set()]
for role, qtext, qimgs in messages:
if role in round_roles[-1]:
# a new round (not the first round)
input_ids += tokenizer(
f"\n[Round {len(round_roles)}]\n{role}：",
add_special_tokens=False,
).input_ids
round_roles.append({role})
else:
round_roles[-1].add(role)
input_ids += tokenizer(
# For first role, no new line
f"\n{role}：" if len(input_ids) != 0 else f"{role}：", add_special_tokens=False
).input_ids
cur_index = 0
for qimg, img_idx in qimgs:
if img_idx > cur_index:
input_ids += tokenizer(
qtext[cur_index:img_idx], add_special_tokens=False
).input_ids
cur_index = img_idx
# image slot, embedding will be replaced by image embeddings
image_slots.append(len(input_ids))
input_ids += [tokenizer.unk_token_id] * nvtokens
images.append(qimg)
input_ids += tokenizer(
qtext[cur_index:], add_special_tokens=False
).input_ids
if len(round_roles) == 1:
# only 1 round
if len(round_roles[0]) == 1 and user_role in round_roles[0]:
# only user role
input_ids += tokenizer("").input_ids
else:
input_ids += tokenizer(f"\n{bot_role}：").input_ids
else:
# add tag for round 0
input_ids = (
tokenizer(f"[Round 0]\n", add_special_tokens=False).input_ids
+ input_ids
)
input_ids += tokenizer(f"\n{bot_role}：").input_ids
if len(input_ids) >= max_length:
image_slots_after_truncate = []
images_after_truncate = []
truncate_index = len(input_ids) - max_length
for image_slot, image in zip(image_slots, images):
# truncate from left
if len(input_ids) - image_slot < max_length:
image_slots_after_truncate.append(image_slot)
images_after_truncate.append(image)
elif len(input_ids) - (image_slot + nvtokens) < max_length:
# in-contact image slot is not allowed
truncate_index = max(truncate_index, image_slot + nvtokens)
for i, image_slot in enumerate(image_slots_after_truncate):
image_slots_after_truncate[i] = image_slot - truncate_index
input_ids = input_ids[truncate_index:]
image_slots = image_slots_after_truncate
images = images_after_truncate
# print(tokenizer.convert_ids_to_tokens(input_ids))
all_images.extend(images)
all_image_slots.append(image_slots)
all_input_ids.append(input_ids)
# 2. Prepare image embeddings
if len(all_images) != 0:
vision_outputs = self.vision_model.forward(torch.cat(all_images, dim=0))
all_image_embeds = vision_outputs[0]
indices_or_sections = [len(chunk) for chunk in all_image_slots]
indices_or_sections = np.cumsum(indices_or_sections)
all_vtokens = []
# TODO: qformer not batched
for image_embeds in torch.tensor_split(
all_image_embeds, tuple(indices_or_sections)
):
image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(
device
)
query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
query_outputs = self.qformer.forward(
query_embeds=query_tokens,
encoder_hidden_states=image_embeds,
encoder_attention_mask=image_atts,
)
query_output = query_outputs[0]
all_vtokens.append(self.language_projection(query_output))
else:
all_vtokens = None
# 3. Place image embeddings into slots
input_ids = (
torch.ones(
(len(all_input_ids), max(len(ids) for ids in all_input_ids)),
dtype=torch.long,
)
* tokenizer.pad_token_id
)
for i, ids in enumerate(all_input_ids):
# pad left
input_ids[i][-len(ids) :] = torch.as_tensor(ids, dtype=torch.long)
input_ids = input_ids.to(device)
inputs_embeds = self.language_model.transformer.word_embeddings(input_ids)
if all_vtokens is not None:
for i, (image_slots, vtokens) in enumerate(
zip(all_image_slots, all_vtokens)
):
for slot, vimg in zip(image_slots, vtokens):
inputs_embeds[i][slot : slot + nvtokens, :] = vimg
return input_ids, inputs_embeds
@torch.no_grad()
def batch_chat(
self,
tokenizer: PreTrainedTokenizer,
batch_messages: List[List[Tuple[str, str, List[Tuple[torch.Tensor, int]]]]],
max_length: int = 2048,
num_beams=1,
do_sample=True,
top_p=0.7,
temperature=0.95,
user_role: str = "问",
bot_role: str = "答",
**kwargs,
):
input_ids, inputs_embeds = self.prepare_inputs_for_chat(
tokenizer=tokenizer,
batch_messages=batch_messages,
max_length=max_length,
user_role=user_role,
bot_role=bot_role,
)
logits_processor = LogitsProcessorList()
logits_processor.append(InvalidScoreLogitsProcessor())
gen_kwargs = {
"max_length": max_length,
"num_beams": num_beams,
"do_sample": do_sample,
"top_p": top_p,
"temperature": temperature,
"logits_processor": logits_processor,
**kwargs,
}
outputs = self.language_model.generate(
input_ids=input_ids, inputs_embeds=inputs_embeds, **gen_kwargs
)
responses = []
for i, output in enumerate(outputs.tolist()):
output = output[len(input_ids[i]) :]
response = tokenizer.decode(output)
responses.append(self.language_model.process_response(response))
return responses
@torch.no_grad()
def stream_chat(
self,
tokenizer: PreTrainedTokenizer,
messages: List[Tuple[str, str, List[Tuple[torch.Tensor, int]]]],
num_beams=5,
max_length=512,
top_p=0.9,
do_sample=True,
temperature=1,
user_role: str = "问",
bot_role: str = "答",
**kwargs,
):
input_ids, inputs_embeds = self.prepare_inputs_for_chat(
tokenizer=tokenizer,
batch_messages=[messages],
max_length=max_length,
user_role=user_role,
bot_role=bot_role,
)
logits_processor = LogitsProcessorList()
logits_processor.append(InvalidScoreLogitsProcessor())
gen_kwargs = {
"max_length": max_length,
"num_beams": num_beams,
"do_sample": do_sample,
"top_p": top_p,
"temperature": temperature,
"logits_processor": logits_processor,
**kwargs,
}
for outputs in self.language_model.stream_generate(
input_ids=input_ids, inputs_embeds=inputs_embeds, **gen_kwargs
):
outputs = outputs.tolist()[0][len(input_ids[0]) :]
response = tokenizer.decode(outputs)
response = self.language_model.process_response(response)
yield response