Spaces:
Running
on
Zero
Running
on
Zero
Delete builder.py
Browse files- builder.py +0 -167
builder.py
DELETED
@@ -1,167 +0,0 @@
|
|
1 |
-
# Copyright 2023 Haotian Liu
|
2 |
-
#
|
3 |
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
4 |
-
# you may not use this file except in compliance with the License.
|
5 |
-
# You may obtain a copy of the License at
|
6 |
-
#
|
7 |
-
# http://www.apache.org/licenses/LICENSE-2.0
|
8 |
-
#
|
9 |
-
# Unless required by applicable law or agreed to in writing, software
|
10 |
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
11 |
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12 |
-
# See the License for the specific language governing permissions and
|
13 |
-
# limitations under the License.
|
14 |
-
|
15 |
-
|
16 |
-
import os
|
17 |
-
import warnings
|
18 |
-
import shutil
|
19 |
-
|
20 |
-
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
|
21 |
-
import torch
|
22 |
-
from llava.model import *
|
23 |
-
from llava.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
|
24 |
-
|
25 |
-
|
26 |
-
def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda", use_flash_attn=False, **kwargs):
|
27 |
-
kwargs = {"device_map": device_map, **kwargs}
|
28 |
-
|
29 |
-
if device != "cuda":
|
30 |
-
kwargs['device_map'] = {"": device}
|
31 |
-
|
32 |
-
if load_8bit:
|
33 |
-
kwargs['load_in_8bit'] = True
|
34 |
-
elif load_4bit:
|
35 |
-
kwargs['load_in_4bit'] = True
|
36 |
-
kwargs['quantization_config'] = BitsAndBytesConfig(
|
37 |
-
load_in_4bit=True,
|
38 |
-
bnb_4bit_compute_dtype=torch.float16,
|
39 |
-
bnb_4bit_use_double_quant=True,
|
40 |
-
bnb_4bit_quant_type='nf4'
|
41 |
-
)
|
42 |
-
else:
|
43 |
-
kwargs['torch_dtype'] = torch.float16
|
44 |
-
|
45 |
-
if use_flash_attn:
|
46 |
-
kwargs['attn_implementation'] = 'flash_attention_2'
|
47 |
-
token = kwargs.get('token', None)
|
48 |
-
if 'llava' in model_name.lower():
|
49 |
-
# Load LLaVA model
|
50 |
-
if 'lora' in model_name.lower() and model_base is None:
|
51 |
-
warnings.warn('There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged.')
|
52 |
-
if 'lora' in model_name.lower() and model_base is not None:
|
53 |
-
from llava.model.language_model.llava_llama import LlavaConfig
|
54 |
-
lora_cfg_pretrained = LlavaConfig.from_pretrained(model_path)
|
55 |
-
tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
|
56 |
-
print('Loading LLaVA from base model...')
|
57 |
-
model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs)
|
58 |
-
token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
|
59 |
-
if model.lm_head.weight.shape[0] != token_num:
|
60 |
-
model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
|
61 |
-
model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
|
62 |
-
|
63 |
-
print('Loading additional LLaVA weights...')
|
64 |
-
if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')):
|
65 |
-
non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), map_location='cpu')
|
66 |
-
else:
|
67 |
-
# this is probably from HF Hub
|
68 |
-
from huggingface_hub import hf_hub_download
|
69 |
-
def load_from_hf(repo_id, filename, subfolder=None):
|
70 |
-
cache_file = hf_hub_download(
|
71 |
-
repo_id=repo_id,
|
72 |
-
filename=filename,
|
73 |
-
subfolder=subfolder)
|
74 |
-
return torch.load(cache_file, map_location='cpu')
|
75 |
-
non_lora_trainables = load_from_hf(model_path, 'non_lora_trainables.bin')
|
76 |
-
non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()}
|
77 |
-
if any(k.startswith('model.model.') for k in non_lora_trainables):
|
78 |
-
non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()}
|
79 |
-
model.load_state_dict(non_lora_trainables, strict=False)
|
80 |
-
|
81 |
-
from peft import PeftModel
|
82 |
-
print('Loading LoRA weights...')
|
83 |
-
model = PeftModel.from_pretrained(model, model_path)
|
84 |
-
print('Merging LoRA weights...')
|
85 |
-
model = model.merge_and_unload()
|
86 |
-
print('Model is loaded...')
|
87 |
-
elif model_base is not None:
|
88 |
-
# this may be mm projector only
|
89 |
-
print('Loading LLaVA from base model...')
|
90 |
-
if 'mpt' in model_name.lower():
|
91 |
-
if not os.path.isfile(os.path.join(model_path, 'configuration_mpt.py')):
|
92 |
-
shutil.copyfile(os.path.join(model_base, 'configuration_mpt.py'), os.path.join(model_path, 'configuration_mpt.py'))
|
93 |
-
tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=True)
|
94 |
-
cfg_pretrained = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
95 |
-
model = LlavaMptForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
|
96 |
-
else:
|
97 |
-
tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
|
98 |
-
cfg_pretrained = AutoConfig.from_pretrained(model_path)
|
99 |
-
model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
|
100 |
-
|
101 |
-
mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
|
102 |
-
mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
|
103 |
-
model.load_state_dict(mm_projector_weights, strict=False)
|
104 |
-
else:
|
105 |
-
if 'mpt' in model_name.lower():
|
106 |
-
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
|
107 |
-
model = LlavaMptForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
|
108 |
-
elif 'mistral' in model_name.lower():
|
109 |
-
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
110 |
-
model = LlavaMistralForCausalLM.from_pretrained(
|
111 |
-
model_path,
|
112 |
-
low_cpu_mem_usage=True,
|
113 |
-
**kwargs
|
114 |
-
)
|
115 |
-
else:
|
116 |
-
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, token=token)
|
117 |
-
model = LlavaLlamaForCausalLM.from_pretrained(
|
118 |
-
model_path,
|
119 |
-
low_cpu_mem_usage=True,
|
120 |
-
**kwargs
|
121 |
-
)
|
122 |
-
else:
|
123 |
-
# Load language model
|
124 |
-
if model_base is not None:
|
125 |
-
# PEFT model
|
126 |
-
from peft import PeftModel
|
127 |
-
tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
|
128 |
-
model = AutoModelForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, **kwargs)
|
129 |
-
print(f"Loading LoRA weights from {model_path}")
|
130 |
-
model = PeftModel.from_pretrained(model, model_path)
|
131 |
-
print(f"Merging weights")
|
132 |
-
model = model.merge_and_unload()
|
133 |
-
print('Convert to FP16...')
|
134 |
-
model.to(torch.float16)
|
135 |
-
else:
|
136 |
-
use_fast = False
|
137 |
-
if 'mpt' in model_name.lower():
|
138 |
-
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
|
139 |
-
model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs)
|
140 |
-
else:
|
141 |
-
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
|
142 |
-
model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
|
143 |
-
|
144 |
-
image_processor = None
|
145 |
-
|
146 |
-
if 'llava' in model_name.lower():
|
147 |
-
mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
|
148 |
-
mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
|
149 |
-
if mm_use_im_patch_token:
|
150 |
-
tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
|
151 |
-
if mm_use_im_start_end:
|
152 |
-
tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
|
153 |
-
model.resize_token_embeddings(len(tokenizer))
|
154 |
-
|
155 |
-
vision_tower = model.get_vision_tower()
|
156 |
-
if not vision_tower.is_loaded:
|
157 |
-
vision_tower.load_model(device_map=device_map)
|
158 |
-
if device_map != 'auto':
|
159 |
-
vision_tower.to(device=device_map, dtype=torch.float16)
|
160 |
-
image_processor = vision_tower.image_processor
|
161 |
-
|
162 |
-
if hasattr(model.config, "max_sequence_length"):
|
163 |
-
context_len = model.config.max_sequence_length
|
164 |
-
else:
|
165 |
-
context_len = 2048
|
166 |
-
|
167 |
-
return tokenizer, model, image_processor, context_len
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|