forestcalled/text-generation-webui

1# Copied from https://github.com/johnsmith0031/alpaca_lora_4bit2 3from pathlib import Path4 5import alpaca_lora_4bit.autograd_4bit as autograd_4bit6from alpaca_lora_4bit.amp_wrapper import AMPWrapper7from alpaca_lora_4bit.autograd_4bit import (8    Autograd4bitQuantLinear,9    load_llama_model_4bit_low_ram10)11from alpaca_lora_4bit.models import Linear4bitLt12from alpaca_lora_4bit.monkeypatch.peft_tuners_lora_monkey_patch import (13    replace_peft_model_with_int4_lora_model14)15 16from modules import shared17from modules.GPTQ_loader import find_quantized_model_file18 19replace_peft_model_with_int4_lora_model()20 21 22def load_model_llama(model_name):23    config_path = str(Path(f'{shared.args.model_dir}/{model_name}'))24    model_path = str(find_quantized_model_file(model_name))25    model, tokenizer = load_llama_model_4bit_low_ram(config_path, model_path, groupsize=shared.args.groupsize, is_v1_model=False)26    for _, m in model.named_modules():27        if isinstance(m, Autograd4bitQuantLinear) or isinstance(m, Linear4bitLt):28            if m.is_v1_model:29                m.zeros = m.zeros.half()30            m.scales = m.scales.half()31            m.bias = m.bias.half()32 33    autograd_4bit.auto_switch = True34 35    model.half()36    wrapper = AMPWrapper(model)37    wrapper.apply_generate()38 39    return model, tokenizer40