Create handler.py

handler.py

@@ -0,0 +1,57 @@
+from typing import Dict, List, Any
+from unsloth.chat_templates import get_chat_template
+from unsloth import FastLanguageModel
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Preload all the elements you are going to need at inference.
+        # pseudo:
+        # self.model= load_model(path)
+        max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
+        dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+        load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name = path, # YOUR MODEL YOU USED FOR TRAINING
+            max_seq_length = max_seq_length,
+            dtype = dtype,
+            load_in_4bit = load_in_4bit,
+            # token=hftoken
+        )
+        FastLanguageModel.for_inference(model)
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `str` | `PIL.Image` | `np.array`)
+            kwargs
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+
+        # pseudo
+        # self.model(input)
+        messages = data
+        # tokenizer = self.tokenizer
+        self.tokenizer = get_chat_template(
+            self.tokenizer,
+            chat_template = "chatml", # Supports zephyr, chatml, mistral, llama, alpaca, vicuna, vicuna_old, unsloth
+            mapping = {"role" : "from",
+                      "content" : "value",
+                      "user" : "human",
+                      "assistant" : "gpt"}, # ShareGPT style
+            map_eos_token = True, # Maps <|im_end|> to  instead
+        )
+        inputs = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize = True,
+            add_generation_prompt = True, # Must add for generation
+            return_tensors = "pt",
+        ).to("cuda")
+
+        # from transformers import TextStreamer
+        # text_streamer = TextStreamer(tokenizer)
+        # _ = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128, use_cache = True)
+        outputs = self.model.generate(input_ids = inputs, max_new_tokens = 64, use_cache = True)
+        # print(outputs)
+        return self.tokenizer.batch_decode(outputs)