import os import sys import gc import torch from transformers import LlamaForCausalLM, LlamaTokenizer from peft import PeftModel from .globals import Global def get_device(): if torch.cuda.is_available(): return "cuda" else: return "cpu" try: if torch.backends.mps.is_available(): return "mps" except: # noqa: E722 pass def get_new_base_model(base_model_name): if Global.ui_dev_mode: return device = get_device() if device == "cuda": model = LlamaForCausalLM.from_pretrained( base_model_name, load_in_8bit=Global.load_8bit, torch_dtype=torch.float16, # device_map="auto", device_map={'': 0}, # ? https://github.com/tloen/alpaca-lora/issues/21 ) elif device == "mps": model = LlamaForCausalLM.from_pretrained( base_model_name, device_map={"": device}, torch_dtype=torch.float16, ) else: model = LlamaForCausalLM.from_pretrained( base_model_name, device_map={"": device}, low_cpu_mem_usage=True ) model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0 model.config.bos_token_id = 1 model.config.eos_token_id = 2 return model def get_tokenizer(base_model_name): if Global.ui_dev_mode: return loaded_tokenizer = Global.loaded_tokenizers.get(base_model_name) if loaded_tokenizer: return loaded_tokenizer tokenizer = LlamaTokenizer.from_pretrained(base_model_name) Global.loaded_tokenizers.set(base_model_name, tokenizer) return tokenizer def get_model( base_model_name, peft_model_name = None): if Global.ui_dev_mode: return if peft_model_name == "None": peft_model_name = None model_key = base_model_name if peft_model_name: model_key = f"{base_model_name}//{peft_model_name}" loaded_model = Global.loaded_models.get(model_key) if loaded_model: return loaded_model peft_model_name_or_path = peft_model_name lora_models_directory_path = os.path.join(Global.data_dir, "lora_models") possible_lora_model_path = os.path.join(lora_models_directory_path, peft_model_name) if os.path.isdir(possible_lora_model_path): peft_model_name_or_path = possible_lora_model_path Global.loaded_models.prepare_to_set() clear_cache() model = get_new_base_model(base_model_name) if peft_model_name: device = get_device() if device == "cuda": model = PeftModel.from_pretrained( model, peft_model_name_or_path, torch_dtype=torch.float16, device_map={'': 0}, # ? https://github.com/tloen/alpaca-lora/issues/21 ) elif device == "mps": model = PeftModel.from_pretrained( model, peft_model_name_or_path, device_map={"": device}, torch_dtype=torch.float16, ) else: model = PeftModel.from_pretrained( model, peft_model_name_or_path, device_map={"": device}, ) model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0 model.config.bos_token_id = 1 model.config.eos_token_id = 2 if not Global.load_8bit: model.half() # seems to fix bugs for some users. model.eval() if torch.__version__ >= "2" and sys.platform != "win32": model = torch.compile(model) Global.loaded_models.set(model_key, model) clear_cache() return model def clear_cache(): gc.collect() # if not shared.args.cpu: # will not be running on CPUs anyway with torch.no_grad(): torch.cuda.empty_cache() def unload_models(): Global.loaded_models.clear() Global.loaded_tokenizers.clear() clear_cache() ######## # def get_base_model(): # load_base_model() # return Global.loaded_base_model # def get_model_with_lora(lora_weights_name_or_path: str = "tloen/alpaca-lora-7b"): # # Global.model_has_been_used = True # # # # # if Global.loaded_tokenizer is None: # Global.loaded_tokenizer = LlamaTokenizer.from_pretrained( # Global.base_model # ) # if Global.cached_lora_models: # model_from_cache = Global.cached_lora_models.get(lora_weights_name_or_path) # if model_from_cache: # return model_from_cache # Global.cached_lora_models.prepare_to_set() # if device == "cuda": # model = PeftModel.from_pretrained( # get_new_base_model(), # lora_weights_name_or_path, # torch_dtype=torch.float16, # device_map={'': 0}, # ? https://github.com/tloen/alpaca-lora/issues/21 # ) # elif device == "mps": # model = PeftModel.from_pretrained( # get_new_base_model(), # lora_weights_name_or_path, # device_map={"": device}, # torch_dtype=torch.float16, # ) # else: # model = PeftModel.from_pretrained( # get_new_base_model(), # lora_weights_name_or_path, # device_map={"": device}, # ) # model.config.pad_token_id = get_tokenizer().pad_token_id = 0 # model.config.bos_token_id = 1 # model.config.eos_token_id = 2 # if not Global.load_8bit: # model.half() # seems to fix bugs for some users. # model.eval() # if torch.__version__ >= "2" and sys.platform != "win32": # model = torch.compile(model) # if Global.cached_lora_models: # Global.cached_lora_models.set(lora_weights_name_or_path, model) # clear_cache() # return model # def load_base_model(): # return; # if Global.ui_dev_mode: # return # if Global.loaded_tokenizer is None: # Global.loaded_tokenizer = LlamaTokenizer.from_pretrained( # Global.base_model # ) # if Global.loaded_base_model is None: # if device == "cuda": # Global.loaded_base_model = LlamaForCausalLM.from_pretrained( # Global.base_model, # load_in_8bit=Global.load_8bit, # torch_dtype=torch.float16, # # device_map="auto", # device_map={'': 0}, # ? https://github.com/tloen/alpaca-lora/issues/21 # ) # elif device == "mps": # Global.loaded_base_model = LlamaForCausalLM.from_pretrained( # Global.base_model, # device_map={"": device}, # torch_dtype=torch.float16, # ) # else: # Global.loaded_base_model = LlamaForCausalLM.from_pretrained( # Global.base_model, device_map={"": device}, low_cpu_mem_usage=True # ) # Global.loaded_base_model.config.pad_token_id = get_tokenizer().pad_token_id = 0 # Global.loaded_base_model.config.bos_token_id = 1 # Global.loaded_base_model.config.eos_token_id = 2 # def clear_cache(): # gc.collect() # # if not shared.args.cpu: # will not be running on CPUs anyway # with torch.no_grad(): # torch.cuda.empty_cache() # def unload_models(): # del Global.loaded_base_model # Global.loaded_base_model = None # del Global.loaded_tokenizer # Global.loaded_tokenizer = None # Global.cached_lora_models.clear() # clear_cache() # Global.model_has_been_used = False # def unload_models_if_already_used(): # if Global.model_has_been_used: # unload_models()