RaDialog-interactive-radiology-report-generation
/
LLAVA_Biovil
/llava
/model
/language_model
/llava_llama.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. | |
import time | |
from packaging import version | |
from typing import List, Optional, Tuple, Union | |
import time | |
import torch | |
import torch.nn as nn | |
import transformers | |
from transformers import AutoConfig, AutoModelForCausalLM, \ | |
LlamaConfig, LlamaModel, LlamaForCausalLM | |
from transformers.modeling_outputs import CausalLMOutputWithPast | |
from LLAVA_Biovil.llava.model.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM | |
class LlavaConfig(LlamaConfig): | |
model_type = "llava" | |
class LlavaLlamaModel(LlavaMetaModel, LlamaModel): | |
config_class = LlavaConfig | |
def __init__(self, config: LlamaConfig, mv_type='none'): | |
super(LlavaLlamaModel, self).__init__(config, mv_type=mv_type) | |
class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM): | |
config_class = LlavaConfig | |
def __init__(self, config, mv_type='none'): | |
super(LlamaForCausalLM, self).__init__(config) | |
self.model = LlavaLlamaModel(config, mv_type=mv_type) | |
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.mv_type = mv_type | |
# 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, | |
prev_images: Optional[torch.FloatTensor] = None, | |
return_dict: Optional[bool] = 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, | |
prev_images | |
) | |
output = 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 | |
) | |
output['modified_labels'] = labels | |
return output | |
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 | |
if version.parse(transformers.__version__) >= version.parse("4.35.0"): | |
AutoConfig.register("llava", LlavaConfig, exist_ok=True) | |
AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM, exist_ok=True) | |
else: | |
AutoConfig.register("llava", LlavaConfig) | |
AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM) | |