Create handler.py

handler.py

@@ -0,0 +1,83 @@
+import torch
+import transformers
+import quant
+from typing import Dict, Any
+from gptq import GPTQ
+from utils import find_layers, DEV
+from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
+
+class EndpointHandler:
+    def __init__(self, 
+                 model_name="Wizard-Vicuna-13B-Uncensored-GPTQ", 
+                 checkpoint_path="Wizard-Vicuna-13B-Uncensored-GPTQ/Wizard-Vicuna-13B-Uncensored-GPTQ-4bit-128g.compat.no-act-order.safetensors",
+                 wbits = 4,
+                 groupsize=128, 
+                 fused_mlp=True, 
+                 eval=True, 
+                 warmup_autotune=True):
+        
+        self.model = self.load_quant(model_name, checkpoint_path, wbits, groupsize, fused_mlp, eval, warmup_autotune)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+        self.model.to(DEV)
+
+    def load_quant(self, model, checkpoint, wbits, groupsize, fused_mlp, eval, warmup_autotune):
+        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
+        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 __call__(self, data: Any) -> Dict[str, str]:
+        input_text = data.pop("inputs", data)
+        input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to(DEV)
+
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                input_ids,
+                do_sample=True,
+                min_length=50,
+                max_length=200,
+                top_p=0.95,
+                temperature=0.8,
+            )
+        generated_text = self.tokenizer.decode([el.item() for el in generated_ids[0]])
+
+        return {'generated_text': generated_text}