dwolfe66/text-generation-webui-space

1import json2import os3import time4import zipfile5from pathlib import Path6 7import numpy as np8import torch9import transformers10from transformers import AutoModelForCausalLM, AutoTokenizer11 12import modules.shared as shared13 14transformers.logging.set_verbosity_error()15 16local_rank = None17 18if shared.args.flexgen:19    from flexgen.flex_opt import (CompressionConfig, ExecutionEnv, OptLM,20                                  Policy, str2bool)21 22if shared.args.deepspeed:23    import deepspeed24    from transformers.deepspeed import (HfDeepSpeedConfig,25                                        is_deepspeed_zero3_enabled)26 27    from modules.deepspeed_parameters import generate_ds_config28 29    # Distributed setup30    local_rank = shared.args.local_rank if shared.args.local_rank is not None else int(os.getenv("LOCAL_RANK", "0"))31    world_size = int(os.getenv("WORLD_SIZE", "1"))32    torch.cuda.set_device(local_rank)33    deepspeed.init_distributed()34    ds_config = generate_ds_config(shared.args.bf16, 1 * world_size, shared.args.nvme_offload_dir)35    dschf = HfDeepSpeedConfig(ds_config) # Keep this object alive for the Transformers integration36 37 38def load_model(model_name):39    print(f"Loading {model_name}...")40    t0 = time.time()41 42    shared.is_RWKV = model_name.lower().startswith('rwkv-')43 44    # Default settings45    if not any([shared.args.cpu, shared.args.load_in_8bit, shared.args.gptq_bits, shared.args.auto_devices, shared.args.disk, shared.args.gpu_memory is not None, shared.args.cpu_memory is not None, shared.args.deepspeed, shared.args.flexgen, shared.is_RWKV]):46        if any(size in shared.model_name.lower() for size in ('13b', '20b', '30b')):47            model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), device_map='auto', load_in_8bit=True)48        else:49            model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16).cuda()50 51    # FlexGen52    elif shared.args.flexgen:53        # Initialize environment54        env = ExecutionEnv.create(shared.args.disk_cache_dir)55 56        # Offloading policy57        policy = Policy(1, 1,58                        shared.args.percent[0], shared.args.percent[1],59                        shared.args.percent[2], shared.args.percent[3],60                        shared.args.percent[4], shared.args.percent[5],61                        overlap=True, sep_layer=True, pin_weight=shared.args.pin_weight,62                        cpu_cache_compute=False, attn_sparsity=1.0,63                        compress_weight=shared.args.compress_weight,64                        comp_weight_config=CompressionConfig(65                            num_bits=4, group_size=64,66                            group_dim=0, symmetric=False),67                        compress_cache=False,68                        comp_cache_config=CompressionConfig(69                            num_bits=4, group_size=64,70                            group_dim=2, symmetric=False))71 72        model = OptLM(f"facebook/{shared.model_name}", env, "models", policy)73 74    # DeepSpeed ZeRO-375    elif shared.args.deepspeed:76        model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), torch_dtype=torch.bfloat16 if shared.args.bf16 else torch.float16)77        model = deepspeed.initialize(model=model, config_params=ds_config, model_parameters=None, optimizer=None, lr_scheduler=None)[0]78        model.module.eval() # Inference79        print(f"DeepSpeed ZeRO-3 is enabled: {is_deepspeed_zero3_enabled()}")80 81    # RMKV model (not on HuggingFace)82    elif shared.is_RWKV:83        from modules.RWKV import RWKVModel, RWKVTokenizer84 85        model = RWKVModel.from_pretrained(Path(f'models/{model_name}'), dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16", device="cpu" if shared.args.cpu else "cuda")86        tokenizer = RWKVTokenizer.from_pretrained(Path('models'))87 88        return model, tokenizer89 90    # Quantized model91    elif shared.args.gptq_bits > 0:92        from modules.GPTQ_loader import load_quantized93 94        model = load_quantized(model_name)95 96    # Custom97    else:98        command = "AutoModelForCausalLM.from_pretrained"99        params = ["low_cpu_mem_usage=True"]100        if not shared.args.cpu and not torch.cuda.is_available():101            print("Warning: no GPU has been detected.\nFalling back to CPU mode.\n")102            shared.args.cpu = True103 104        if shared.args.cpu:105            params.append("low_cpu_mem_usage=True")106            params.append("torch_dtype=torch.float32")107        else:108            params.append("device_map='auto'")109            params.append("load_in_8bit=True" if shared.args.load_in_8bit else "torch_dtype=torch.bfloat16" if shared.args.bf16 else "torch_dtype=torch.float16")110 111            if shared.args.gpu_memory:112                memory_map = shared.args.gpu_memory113                max_memory = f"max_memory={{0: '{memory_map[0]}GiB'"114                for i in range(1, len(memory_map)):115                    max_memory += (f", {i}: '{memory_map[i]}GiB'")116                max_memory += (f", 'cpu': '{shared.args.cpu_memory or '99'}GiB'}}")117                params.append(max_memory)118            elif not shared.args.load_in_8bit:119                total_mem = (torch.cuda.get_device_properties(0).total_memory/(1024*1024))120                suggestion = round((total_mem-1000)/1000)*1000121                if total_mem-suggestion < 800:122                    suggestion -= 1000123                suggestion = int(round(suggestion/1000))124                print(f"\033[1;32;1mAuto-assiging --gpu-memory {suggestion} for your GPU to try to prevent out-of-memory errors.\nYou can manually set other values.\033[0;37;0m")125                params.append(f"max_memory={{0: '{suggestion}GiB', 'cpu': '{shared.args.cpu_memory or '99'}GiB'}}")126            if shared.args.disk:127                params.append(f"offload_folder='{shared.args.disk_cache_dir}'")128 129        command = f"{command}(Path(f'models/{shared.model_name}'), {', '.join(set(params))})"130        model = eval(command)131 132    # Loading the tokenizer133    if shared.model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')) and Path("models/gpt-j-6B/").exists():134        tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/"))135    else:136        tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{shared.model_name}/"))137    tokenizer.truncation_side = 'left'138 139    print(f"Loaded the model in {(time.time()-t0):.2f} seconds.")140    return model, tokenizer141 142def load_soft_prompt(name):143    if name == 'None':144        shared.soft_prompt = False145        shared.soft_prompt_tensor = None146    else:147        with zipfile.ZipFile(Path(f'softprompts/{name}.zip')) as zf:148            zf.extract('tensor.npy')149            zf.extract('meta.json')150            j = json.loads(open('meta.json', 'r').read())151            print(f"\nLoading the softprompt \"{name}\".")152            for field in j:153                if field != 'name':154                    if type(j[field]) is list:155                        print(f"{field}: {', '.join(j[field])}")156                    else:157                        print(f"{field}: {j[field]}")158            print()159            tensor = np.load('tensor.npy')160            Path('tensor.npy').unlink()161            Path('meta.json').unlink()162        tensor = torch.Tensor(tensor).to(device=shared.model.device, dtype=shared.model.dtype)163        tensor = torch.reshape(tensor, (1, tensor.shape[0], tensor.shape[1]))164 165        shared.soft_prompt = True166        shared.soft_prompt_tensor = tensor167 168    return name169