|
|
import os |
|
|
from modeling_videochat2 import * |
|
|
from modeling_base import freeze_module |
|
|
from transformers import AutoConfig |
|
|
token = os.environ['HF_TOKEN'] |
|
|
|
|
|
|
|
|
class InternVideo2_cls(InternVideo2_VideoChat2): |
|
|
def __init__(self, config): |
|
|
super(InternVideo2_VideoChat2, self).__init__(config=config) |
|
|
|
|
|
def build_llm(self): |
|
|
self.lm_name = self.model_config.llm.name |
|
|
if self.model_config.llm.name == 'mistral_7b': |
|
|
from transformers import AutoModelForSequenceClassification |
|
|
config = AutoConfig.from_pretrained( |
|
|
self.model_config.llm.pretrained_llm_path, |
|
|
torch_dtype=torch.bfloat16, |
|
|
token=token, |
|
|
num_labels=self.model_config.llm.num_labels |
|
|
|
|
|
) |
|
|
self.lm = AutoModelForSequenceClassification.from_config(config) |
|
|
|
|
|
else: |
|
|
raise NotImplementedError(self.model_config.llm.name) |
|
|
|
|
|
self.freeze_llm = self.model_config.get("freeze_llm", True) |
|
|
logger.info(f'freeze_llm: {self.freeze_llm}') |
|
|
if self.freeze_llm: |
|
|
logger.info("freeze llm") |
|
|
freeze_module(self.lm) |
|
|
|
|
|
if self.model_config.llm.use_lora: |
|
|
self.use_lora = True |
|
|
from peft import get_peft_model, LoraConfig, TaskType |
|
|
logger.info("Use lora") |
|
|
|
|
|
peft_config = LoraConfig( |
|
|
task_type=TaskType.CAUSAL_LM, inference_mode=False, |
|
|
r=self.model_config.llm.lora_r, lora_alpha=self.model_config.llm.lora_alpha, lora_dropout=self.model_config.llm.lora_dropout, |
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", |
|
|
"gate_proj", "up_proj", "down_proj"] |
|
|
) |
|
|
|
|
|
self.lm = get_peft_model(self.lm, peft_config) |
|
|
self.lm.enable_input_require_grads() |
|
|
self.lm.print_trainable_parameters() |
|
|
else: |
|
|
self.use_lora = False |
|
|
|
|
|
|
|
|
def build_conversation(self,instruction, user_prompt,media_type='video',msg=''): |
|
|
|
|
|
conversation = "" |
|
|
if instruction: |
|
|
conversation += instruction |
|
|
conversation += ("[INST]" + " ") |
|
|
|
|
|
if media_type == 'image': |
|
|
conversation +=( "<Image>" + IMG_TOKEN + "</Image>") |
|
|
else: |
|
|
conversation += ("<Video>" + VID_TOKEN + "</Video>") |
|
|
|
|
|
conversation += (msg.rstrip() + "[/INST]") |
|
|
conversation += (" [INST] " + user_prompt + " [/INST]") |
|
|
conversation += ("") |
|
|
return conversation |
|
|
|
|
|
def test(self, x): |
|
|
return x |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained('OpenGVLab/InternVideo2-Chat-8B',trust_remote_code=True,use_fast=False) |
|
|
config = AutoConfig.from_pretrained('OpenGVLab/InternVideo2-Chat-8B', torch_dtype=torch.bfloat16,trust_remote_code=True) |
|
|
model = InternVideo2_Classification(config).cuda() |
|
|
|
|
|
B, T, C, H, W = 1, 8, 3, 224, 224 |
|
|
video_tensor = torch.randn(B,T,C,H,W).cuda() |
|
|
user_prompt = "this is a user prompt" |
|
|
instruction = "this is an instruction" |
|
|
|
|
|
conversation = model.build_conversation(instruction=instruction, user_prompt=user_prompt, media_type='video') |
|
|
tokenized = model.build_input_ids(tokenizer,conversation,max_length=248,add_special_tokens=True,truncation=False,padding=False,return_tensors='pt') |
|
|
|
|
|
input_ids = tokenized['input_ids'].unsqueeze(0).to(model.device) |
|
|
attn_mask = tokenized['attention_mask'].unsqueeze(0).to(model.device) |
|
|
indexes = tokenized['index'].unsqueeze(0) |
|
|
text_embeds = model.pad_text_embeds(input_ids = input_ids,video = video_tensor,video_idx = indexes) |
|
|
outputs = model.lm(inputs_embeds=text_embeds, attention_mask=attn_mask,output_hidden_states=True,return_dict=True) |
|
|
|
