Create llama_inference_class.py

llama_inference_class.py

@@ -0,0 +1,96 @@
+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]])