File size: 3,037 Bytes
eacf0bd 908a9d5 eacf0bd 9b5bdb0 eacf0bd 908a9d5 eacf0bd 9b5bdb0 eacf0bd 908a9d5 eacf0bd 908a9d5 eacf0bd a3e7589 eacf0bd a3e7589 908a9d5 eacf0bd |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import torch
import warnings
from config import *
from transformers import BitsAndBytesConfig
warnings.filterwarnings(action='ignore')
def setting_trol_config(trol, tok_trol, image_special_token):
trol.config.image_token_index = tok_trol.convert_tokens_to_ids(image_special_token)
trol.config.ignore_index = -100
trol.config.pad_token_id = tok_trol.eos_token_id
trol.config.eos_token_id = tok_trol.eos_token_id
def load_trol(link):
"""
model selection
"""
if link == 'TroL-1.8B':
from .arch_internlm2.modeling_trol import TroLForCausalLM
from .arch_internlm2.tokenization_internlm2 import InternLM2Tokenizer as TroLTokenizer
bits = 4
path = TROL_1_8B
image_special_token = "<image>"
bit_quant_skip = ["vit", "vision_proj", "ffn", "output", "trol_gating"]
elif link == 'TroL-3.8B':
from trol.arch_phi3.modeling_trol import TroLForCausalLM
from transformers import LlamaTokenizerFast as TroLTokenizer
bits = 8
path = TROL_3_8B
image_special_token = "<IMG_CONTEXT>"
bit_quant_skip = ["vision_model", "vision_proj", "lm_head", "trol_gating"]
elif link == 'TroL-7B':
from .arch_internlm2.modeling_trol import TroLForCausalLM
from .arch_internlm2.tokenization_internlm2 import InternLM2Tokenizer as TroLTokenizer
bits = 4
path = TROL_7B
image_special_token = "<image>"
bit_quant_skip = ["vit", "vision_proj", "ffn", "output", "trol_gating"]
else:
raise Exception("Unsupported Link")
# huggingface model configuration
huggingface_config = {}
# Bit quantization
if bits in [4, 8]:
huggingface_config.update(dict(
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
attn_implementation="flash_attention_2",
quantization_config=BitsAndBytesConfig(
load_in_4bit=bits == 4,
load_in_8bit=bits == 8,
llm_int8_skip_modules=bit_quant_skip,
llm_int8_threshold=6.0,
llm_int8_has_fp16_weight=False,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type='nf4'
)
))
else:
huggingface_config.update(dict(
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
attn_implementation="flash_attention_2",
))
# Loading tokenizer & Loading backbone model (error -> then delete flash attention)
tok_trol = TroLTokenizer.from_pretrained(path, padding_side='left')
try:
trol = TroLForCausalLM.from_pretrained(path, **huggingface_config)
except:
del huggingface_config["attn_implementation"]
trol = TroLForCausalLM.from_pretrained(path, **huggingface_config)
trol.config.llm_config.use_cache = False
# setting config
setting_trol_config(trol, tok_trol, image_special_token)
return trol, tok_trol |