handler

handler.py

@@ -1,64 +1,28 @@
-# handler.py
-from peft import AutoPeftModelForCausalLM
-from transformers import AutoTokenizer
-from typing import Dict, List, Any
+from typing import Dict, Any, List
+from unsloth import FastLanguageModel
+import torch
+
+max_seq_length = 2048
+dtype = None
+load_in_4bit = True
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # Initialize and load the model using a Transformers pipeline
-        # Ensure that the model and tokenizer are placed correctly in the specified path
-        # self.model = PreTrainedModel.from_pretrained(path).to(device)
-        # self.tokenizer = PreTrainedTokenizer.from_pretrained(path)
-        # max_seq_length = 1024
-        # dtype = None
-        load_in_4bit = True
-        # self.model, self.tokenizer = FastLanguageModel.from_pretrained(
-        #     model_name=path,
-        #     max_seq_length=max_seq_length,
-        #     dtype=dtype,
-        #     load_in_4bit=load_in_4bit
-        # )
-        # FastLanguageModel.for_inference(self.model)
-
-        from transformers import AutoTokenizer
-        self.model = AutoPeftModelForCausalLM.from_pretrained(
-            path,
-            load_in_4bit = load_in_4bit,
+        self.model, self.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
         )
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        FastLanguageModel.for_inference(self.model) # Enable native 2x faster inference
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        """
-        Handles incoming requests to the inference endpoint.
-
-        Args:
-            data (Dict[str, Any]): The request payload containing 'inputs' and any other necessary parameters.
+        inputs = self.tokenizer([
+            """Q:
+            I am having problems with my computer. It is not turning on. What should I do?
+            A:"""
+            ], return_tensors = "pt").to("cuda")
 
-        Returns:
-            List[Dict[str, Any]]: The prediction results formatted as a list of dictionaries.
-        """
-        # encode the input text
-        # inputs = self.tokenizer([
-        #     """<s>\nQ:\nI am having problems with my computer. It is not turning on. What should I do?\n"A:\n"""
-        #     ], return_tensors="pt").to("cuda")
-        # # generate text based on the input
-        # outputs = self.model.generate(**inputs, max_new_tokens=1000, use_cache=True)
-        # result = self.tokenizer.batch_decode(outputs)
-        # return [{"generated_text": result[0]}]
-
-        # encode the input text
-        inputs = self.tokenizer.encode(
-            text="\n<s>\nQ:\nI am having problems with my computer. It is not turning on. What should I do?\nA:\n",
-            return_tensors="pt",
-            max_length=1024,
-            truncation=True,
-            padding=True,
-        ).to(self.model.device)
-        outputs = self.model.generate(
-            inputs,
-            max_length=1024,
-            num_return_sequences=1,
-            use_cache=True,
-        )
-        result = self.tokenizer.batch_decode(outputs)
-        return [{"generated_text": result}]
+        outputs = self.model.generate(**inputs, max_new_tokens = 1000, use_cache = True)
+        results = self.tokenizer.batch_decode(outputs)
+        return results