|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from typing import List, Optional, Tuple, Union |
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
from transformers import AutoConfig, AutoModelForCausalLM, \ |
|
LlamaConfig, LlamaModel, LlamaForCausalLM |
|
from transformers.modeling_outputs import CausalLMOutputWithPast |
|
from transformers.generation.utils import GenerateOutput |
|
|
|
from .videollama2_arch import Videollama2MetaModel, Videollama2MetaForCausalLM |
|
|
|
|
|
class Videollama2LlamaConfig(LlamaConfig): |
|
model_type = "videollama2_llama" |
|
|
|
def __init__(self, **kwargs): |
|
super().__init__(**kwargs) |
|
self.model_type = "videollama2_llama" |
|
|
|
|
|
class Videollama2LlamaModel(Videollama2MetaModel, LlamaModel): |
|
config_class = Videollama2LlamaConfig |
|
|
|
def __init__(self, config: LlamaConfig): |
|
super(Videollama2LlamaModel, self).__init__(config) |
|
|
|
|
|
class Videollama2LlamaForCausalLM(LlamaForCausalLM, Videollama2MetaForCausalLM): |
|
config_class = Videollama2LlamaConfig |
|
|
|
def __init__(self, config, **kwargs): |
|
super(LlamaForCausalLM, self).__init__(config) |
|
self.model = Videollama2LlamaModel(config) |
|
self.pretraining_tp = config.pretraining_tp |
|
self.vocab_size = config.vocab_size |
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
|
|
|
|
|
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, |
|
return_dict: Optional[bool] = None, |
|
cache_position: Optional[torch.LongTensor] = None, |
|
**kwargs |
|
) -> Union[Tuple, CausalLMOutputWithPast]: |
|
|
|
if inputs_embeds is None: |
|
( |
|
input_ids, |
|
attention_mask, |
|
past_key_values, |
|
inputs_embeds, |
|
labels |
|
) = self.prepare_inputs_labels_for_multimodal( |
|
input_ids, |
|
attention_mask, |
|
past_key_values, |
|
labels, |
|
images |
|
) |
|
|
|
outputs = super().forward( |
|
input_ids=input_ids, |
|
attention_mask=attention_mask, |
|
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, |
|
cache_position=cache_position, |
|
) |
|
|
|
outputs.labels = labels |
|
|
|
return outputs |
|
|
|
@torch.no_grad() |
|
def generate( |
|
self, |
|
inputs: Optional[torch.Tensor] = None, |
|
images: 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: |
|
( |
|
input_ids, |
|
attention_mask, |
|
past_key_values, |
|
inputs_embeds, |
|
_ |
|
) = self.prepare_inputs_labels_for_multimodal( |
|
input_ids=inputs, |
|
attention_mask=attention_mask, |
|
past_key_values=None, |
|
labels=None, |
|
images=images |
|
) |
|
else: |
|
inputs_embeds = self.get_model().embed_tokens(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) |
|
_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 |
|
return _inputs |
|
|
|
|
|
AutoConfig.register("videollama2_llama", Videollama2LlamaConfig) |
|
AutoModelForCausalLM.register(Videollama2LlamaConfig, Videollama2LlamaForCausalLM) |
|
|