modeling_blip2chatglm.py

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