File size: 3,196 Bytes
74b17e0 |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import os
import torch
from collections import OrderedDict
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
from .modeling_tinyllava import TinyLlavaForConditionalGeneration
from .configuration_tinyllava import TinyLlavaConfig
def load_base_ckp_for_lora(ckp_path):
ckp = torch.load(ckp_path, map_location=torch.device('cpu'))
new_ckp = OrderedDict()
for k, v in ckp.items():
new_k = k.replace('.base_layer', '')
new_ckp[new_k] = v
return new_ckp
def load_pretrained_model(model_name_or_path, load_type='hf', load_8bit=False, load_4bit=False, device_map="auto",
device="cuda", **kwargs):
kwargs = {"device_map": device_map, **kwargs}
if device != "cuda":
kwargs['device_map'] = {"": device}
if load_8bit:
kwargs['load_in_8bit'] = True
elif load_4bit:
kwargs['load_in_4bit'] = True
kwargs['quantization_config'] = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type='nf4'
)
else:
kwargs['torch_dtype'] = torch.float16
if model_name_or_path is not None and 'lora' not in model_name_or_path:
model = TinyLlavaForConditionalGeneration.from_pretrained(model_name_or_path,low_cpu_mem_usage=True)
elif model_name_or_path is not None and 'lora' in model_name_or_path:
if os.path.exists(os.path.join(model_name_or_path, 'adapter_config.json')):
model_config = TinyLlavaConfig.from_pretrained(model_name_or_path)
model = TinyLlavaForConditionalGeneration(model_config)
language_model_ckp_path = os.path.join(model_name_or_path, 'language_model/pytorch_model.bin')
language_model_ckp = load_base_ckp_for_lora(language_model_ckp_path)
model.language_model.load_state_dict(language_model_ckp)
vision_tower_ckp_path = os.path.join(model_name_or_path, 'vision_tower/pytorch_model.bin')
vision_tower_ckp = load_base_ckp_for_lora(vision_tower_ckp_path)
model.vision_tower._vision_tower.load_state_dict(vision_tower_ckp)
connector_ckp_path = os.path.join(model_name_or_path, 'connector/pytorch_model.bin')
connector_ckp = load_base_ckp_for_lora(connector_ckp_path)
model.connector.load_state_dict(connector_ckp, strict=False)
model.to(torch.float16)
from peft import PeftModel
print('Loading LoRA weights...')
model = PeftModel.from_pretrained(model, model_name_or_path)
print('Merging LoRA weights...')
model = model.merge_and_unload()
print('Model is loaded...')
image_processor = model.vision_tower._image_processor
context_len = getattr(model.config, 'max_sequence_length', 2048)
# tokenizer = AutoTokenizer.from_pretrained(model.config.llm_model_name_or_path, use_fast=False, padding_side="right")
tokenizer = model.tokenizer
#tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer, image_processor, context_len
|