Create handler.py

handler.py

@@ -0,0 +1,61 @@
+from typing import Dict, List, Any
+from peft import PeftModel
+from unsloth import FastLanguageModel
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        llm_model,tokenizer=initialize_model_and_tokenizer("mistralai/Mistral-7B-Instruct-v0.2",2048)
+        llm_model = PeftModel.from_pretrained(llm_model, "./")
+        llm_model.eval()
+        self.llm_model = llm_model
+        self.tokenizer = tokenizer
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+         data args:
+              inputs (:obj: `str`)
+              date (:obj: `str`)
+        Return:
+              A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # get inputs
+        prompt = data.pop("prompt", "")
+        model_input = self.tokenizer(prompt, return_tensors="pt").to(device_map)
+        output = self.llm_model.generate(input_ids=model_input["input_ids"].to(device_map),
+                                  use_cache=False,
+                                  temperature=0.1, top_k=1, top_p=1.0, repetition_penalty=1.4,
+                                  max_new_tokens=256,
+                                  do_sample=True,
+                                  pad_token_id=tokenizer.pad_token_id,
+                                  eos_token_id=tokenizer.eos_token_id,
+                                  num_beams=1,
+                                  num_return_sequences=1)
+        output = self.tokenizer.decode(output[0])
+        result = (output
+                .split(tokenizer.eos_token)[0]
+                .split("Response:")[1]
+                .strip()
+                .split("###")[0]
+                .replace("```json", "")
+                .replace("```", ""))
+
+        return {"response":result}
+
+
+def initialize_model_and_tokenizer(model_id: str, max_seq_length: int):
+    # Initialize model and tokenizer
+    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=model_id,
+        max_seq_length=max_seq_length,
+        dtype=dtype,
+        load_in_4bit=load_in_4bit,
+    )
+    model.config.use_cache = False
+    model.config.pad_token_id = model.config.eos_token_id
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+    return model, tokenizer