File size: 2,086 Bytes
3cd1076 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | import glob
import torch
from transformers import LlamaConfig
from tqdm import tqdm
def update_custom_config(config, model_args):
if model_args.llm_model_name_or_path is not None:
text_config = LlamaConfig.from_pretrained(
model_args.llm_model_name_or_path)
config.text_config = text_config
config.with_qformer = model_args.with_qformer
config.custom_vision_model = model_args.custom_vision_model
if model_args.custom_vision_model:
# config.vision_model_args = model_args
config.image_encoder_name_or_path = \
model_args.image_encoder_name_or_path
config.vision_select_layer = model_args.vision_select_layer
if getattr(model_args, "vision_select_feature", None) is not None:
config.vision_select_feature = model_args.vision_select_feature
return config
def load_llava_pretrain_model(model, checkpoint_path):
checkpoint_path = glob.glob(checkpoint_path)
for path in tqdm(checkpoint_path):
state_dict = torch.load(path, map_location="cpu")
new_state_dict = adapt_llava_model_to_lmflow_type(state_dict)
# modify the name of the key
# import pdb; pdb.set_trace()
lmflow_keys = model.state_dict().keys()
for key in new_state_dict.keys():
if key not in lmflow_keys:
print("key not in lmflow_keys: ", key)
model.load_state_dict(new_state_dict, strict=False)
return model
def adapt_llava_model_to_lmflow_type(state_dict):
new_state_dict = {}
for key, item in state_dict.items():
key = key.replace("model.layers", "language_model.model.layers")
key = key.replace("model.embed_tokens",
"language_model.model.embed_tokens")
key = key.replace("model.mm_projector", "language_projection")
key = key.replace("lm_head", "language_model.lm_head")
key = key.replace("model.norm", "language_model.model.norm")
if "vision_tower" in key:
continue
new_state_dict[key] = item
return new_state_dict |