| import torch | |
| import torch.nn as nn | |
| import quant | |
| from gptq import GPTQ | |
| from utils import find_layers, DEV, set_seed, get_wikitext2, get_ptb, get_c4, get_ptb_new, get_c4_new, get_loaders | |
| import transformers | |
| from transformers import AutoTokenizer | |
| class ModelInference: | |
| def __init__(self, model_name, load=None, wbits=16, groupsize=-1): | |
| self.model_name = model_name | |
| self.load = load | |
| self.wbits = wbits | |
| self.groupsize = groupsize | |
| if self.load: | |
| self.model = self.load_quant(self.model_name, self.load, self.wbits, self.groupsize) | |
| else: | |
| self.model = self.get_llama(self.model_name) | |
| self.model.eval() | |
| self.model.to(DEV) | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False) | |
| def get_llama(model): | |
| def skip(*args, **kwargs): | |
| pass | |
| torch.nn.init.kaiming_uniform_ = skip | |
| torch.nn.init.uniform_ = skip | |
| torch.nn.init.normal_ = skip | |
| from transformers import LlamaForCausalLM | |
| model = LlamaForCausalLM.from_pretrained(model, torch_dtype='auto') | |
| model.seqlen = 2048 | |
| return model | |
| def load_quant(model, checkpoint, wbits, groupsize=-1, fused_mlp=True, eval=True, warmup_autotune=True): | |
| from transformers import LlamaConfig, LlamaForCausalLM | |
| config = LlamaConfig.from_pretrained(model) | |
| def noop(*args, **kwargs): | |
| pass | |
| torch.nn.init.kaiming_uniform_ = noop | |
| torch.nn.init.uniform_ = noop | |
| torch.nn.init.normal_ = noop | |
| torch.set_default_dtype(torch.half) | |
| transformers.modeling_utils._init_weights = False | |
| torch.set_default_dtype(torch.half) | |
| model = LlamaForCausalLM(config) | |
| torch.set_default_dtype(torch.float) | |
| if eval: | |
| model = model.eval() | |
| layers = find_layers(model) | |
| for name in ['lm_head']: | |
| if name in layers: | |
| del layers[name] | |
| quant.make_quant_linear(model, layers, wbits, groupsize) | |
| del layers | |
| print('Loading model ...') | |
| if checkpoint.endswith('.safetensors'): | |
| from safetensors.torch import load_file as safe_load | |
| model.load_state_dict(safe_load(checkpoint), strict=False) | |
| else: | |
| model.load_state_dict(torch.load(checkpoint), strict=False) | |
| if eval: | |
| quant.make_quant_attn(model) | |
| quant.make_quant_norm(model) | |
| if fused_mlp: | |
| quant.make_fused_mlp(model) | |
| if warmup_autotune: | |
| quant.autotune_warmup_linear(model, transpose=not (eval)) | |
| if eval and fused_mlp: | |
| quant.autotune_warmup_fused(model) | |
| model.seqlen = 2048 | |
| print('Done.') | |
| return model | |
| def generate_text(self, text, min_length=10, max_length=50, top_p=0.95, temperature=0.8): | |
| input_ids = self.tokenizer.encode(text, return_tensors="pt").to(DEV) | |
| with torch.no_grad(): | |
| generated_ids = self.model.generate( | |
| input_ids, | |
| do_sample=True, | |
| min_length=min_length, | |
| max_length=max_length, | |
| top_p=top_p, | |
| temperature=temperature, | |
| ) | |
| return self.tokenizer.decode([el.item() for el in generated_ids[0]]) | |