llava_internlm2.py

#    Copyright 2023 Haotian Liu
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss

from transformers import AutoConfig, AutoModelForCausalLM
from llava.model.language_model.internlm2.modeling_internlm2 import InternLM2ForCausalLM, InternLM2Model
from llava.model.language_model.internlm2.configuration_internlm2 import InternLM2Config

from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation.utils import GenerateOutput

from llava.model.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
from llava.utils import rank0_print
# from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM


class LlavaInternlm2Config(InternLM2Config):
    model_type = "llava_internlm2"


class LlavaInternlm2Model(LlavaMetaModel, InternLM2Model):
    config_class = LlavaInternlm2Config

    def __init__(self, config: InternLM2Config):
        super(LlavaInternlm2Model, self).__init__(config)


class LlavaInternlm2ForCausalLM(InternLM2ForCausalLM, LlavaMetaForCausalLM):
    config_class = LlavaInternlm2Config

    def __init__(self, config):
        # super(InternLM2ForCausalLM, self).__init__(config)
        InternLM2ForCausalLM.__init__(self, config)
        self.model = LlavaInternlm2Model(config)
        # self.pretraining_tp = config.pretraining_tp
        self.vocab_size = config.vocab_size
        self.datatype_loss = config.datatype_loss if hasattr(config, "datatype_loss") else False
        if self.datatype_loss:
            rank0_print("Logging per datatype loss")
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_model(self):
        return self.model

    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,
        images: Optional[torch.FloatTensor] = None,
        image_sizes: Optional[List[List[int]]] = None,
        modalities: Optional[List[str]] = ["image"],
        data_type: Optional[str] = "normal",
        return_dict: Optional[bool] = None,
        dpo_forward: Optional[bool] = False,
        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, images, image_sizes
                )
            )
        if not self.datatype_loss:
            if dpo_forward:
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_values=past_key_values,
                    inputs_embeds=inputs_embeds,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=return_dict,
                )

                hidden_states = outputs[0]
                logits = self.output(hidden_states)
                return logits, labels
            else:
                return super().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,
                )
        else:
            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
            )
            return_dict = return_dict if return_dict is not None else self.config.use_return_dict

            # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )

            hidden_states = outputs[0]
            logits = self.output(hidden_states)
            logits = logits.float()

            loss = None
            per_sample_losses = None

            if labels is not None:
                # Shift so that tokens < n predict n
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()
                # Flatten the tokens
                # loss_fct = CrossEntropyLoss()
                shift_logits = shift_logits.view(-1, self.config.vocab_size)
                shift_labels = shift_labels.view(-1)
                # Enable model parallelism
                shift_labels = shift_labels.to(shift_logits.device)
                # loss = loss_fct(shift_logits, shift_labels)

                ##### Compute per sample loss #####
                # Compute the token-level loss
                loss_fct = CrossEntropyLoss(reduction="none")  # "none" for token-level losses
                token_losses = loss_fct(shift_logits, shift_labels)  # Shape: [batch_size * seq_len]

                # Reshape token losses to [batch_size, seq_len - 1]
                token_losses = token_losses.view(-1, shift_logits.size(0) // inputs_embeds.size(0))
                # batch_size = inputs_embeds.size(0)
                # seq_len = inputs_embeds.size(1)
                # token_losses = token_losses.view(batch_size, seq_len - 1)

                # Mask out padding tokens
                active_tokens = (shift_labels != -100).view(-1, token_losses.size(1))
                token_losses *= active_tokens

                # Compute per-sample losses by summing over the sequence length
                per_sample_losses = token_losses.sum(dim=1) / active_tokens.sum(dim=1).clamp(min=1)

                # Compute overall loss as the mean of per-sample losses
                loss = per_sample_losses.mean()

            if not return_dict:
                output = (logits,) + outputs[1:]
                return (loss,) + output if loss is not None else output

            device = input_ids.device if input_ids is not None else inputs_embeds.device
            output = CausalLMOutputWithPast(
                loss=loss,
                logits=logits,
                past_key_values=outputs.past_key_values,
                hidden_states=outputs.hidden_states,
                attentions=outputs.attentions,
            )
            output['logits'] = output['logits'].to(device)

            output["per_sample_losses"] = per_sample_losses  # Include per-sample losses in the output

            return output

    @torch.no_grad()
    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        images: Optional[torch.Tensor] = None,
        image_sizes: Optional[torch.Tensor] = 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 images is not None:
            (inputs, position_ids, attention_mask, _, inputs_embeds, _) = (
                self.prepare_inputs_labels_for_multimodal(
                    inputs, position_ids, attention_mask, None, None, images, image_sizes=image_sizes
                )
            )
        else:
            inputs_embeds = self.get_model().get_input_embeddings()(inputs)

        return super().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):
        images = kwargs.pop("images", None)
        image_sizes = kwargs.pop("image_sizes", None)
        inputs = super().prepare_inputs_for_generation(
            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
        )
        if images is not None:
            inputs['images'] = images
        if image_sizes is not None:
            inputs['image_sizes'] = image_sizes
        return inputs


AutoConfig.register("llava_internlm2", LlavaInternlm2Config)
AutoModelForCausalLM.register(LlavaInternlm2Config, LlavaInternlm2ForCausalLM)