Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,126 Bytes
7c34c28 |
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 84 |
from typing import Type, List, Optional
import logging
from transformers import AutoTokenizer, AutoConfig, BitsAndBytesConfig
from huggingface_hub import hf_hub_download
from peft import PeftModel
import torch
import os
from multi_token.model_utils import fix_tokenizer, MultiTaskType
from multi_token.modalities.base_modality import Modality
from multi_token.language_models.mistral import MistralForCausalLM
from multi_token.language_models import LANGUAGE_MODEL_NAME_TO_CLASS
from multi_token.modalities import MODALITY_BUILDERS
def load_trained_lora_model(
model_name_or_path: str,
model_lora_path: str,
model_cls: Optional[Type] = None,
modalities: Optional[List[Modality]] = None,
load_bits: int = 16,
device_map: str = "auto",
use_multi_task: int = MultiTaskType.NO_MULTI_TASK,
tasks_config: str = None
):
load_kwargs = {"device_map": device_map}
if load_bits == 8:
load_kwargs["load_in_8bit"] = True
elif load_bits == 4:
load_kwargs["load_in_4bit"] = True
load_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",
)
elif load_bits == 16:
load_kwargs["torch_dtype"] = torch.float16
else:
raise ValueError(f"Invalid load_bits: {load_bits}")
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
fix_tokenizer(tokenizer)
cfg = AutoConfig.from_pretrained(model_lora_path)
if model_cls is None:
model_cls = LANGUAGE_MODEL_NAME_TO_CLASS[cfg.model_cls]
if modalities is None:
if use_multi_task:
modalities = MODALITY_BUILDERS[cfg.modality_builder](use_multi_task = use_multi_task, tasks_config = tasks_config)
else:
modalities = MODALITY_BUILDERS[cfg.modality_builder]()
logging.info(f"Loading base model from {model_name_or_path} as {load_bits} bits")
model = model_cls.from_pretrained(
model_name_or_path, low_cpu_mem_usage=True, config=cfg, **load_kwargs
)
model.modalities = modalities
logging.info(f"Loading projector weights for {[m.name for m in modalities]}")
if os.path.exists(os.path.join(model_lora_path, "non_lora_trainables.bin")):
non_lora_trainables = torch.load(
os.path.join(model_lora_path, "non_lora_trainables.bin"), map_location="cuda"
)
else:
local_fn = hf_hub_download(
repo_id=model_lora_path,
filename="non_lora_trainables.bin",
repo_type="model",
)
non_lora_trainables = torch.load(local_fn, map_location="cuda")
model.get_model().initialize_pretrained_modules(modalities, non_lora_trainables)
logging.info(f"Loading and merging LoRA weights from {model_lora_path}")
model = PeftModel.from_pretrained(model, model_lora_path)
if load_bits == 16:
# TODO: Figure out why this fails for other bit sizes
model = model.merge_and_unload()
model.eval()
return model, tokenizer
|