Create handler.py

handler.py

@@ -0,0 +1,63 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import re
+import torch
+
+template = """Alice Gate's Persona: Alice Gate is a young, computer engineer-nerd with a knack for problem solving and a passion for technology.
+<START>
+{user_name}: So how did you get into computer engineering?
+Alice Gate: I've always loved tinkering with technology since I was a kid.
+{user_name}: That's really impressive!
+Alice Gate: *She chuckles bashfully* Thanks!
+{user_name}: So what do you do when you're not working on computers?
+Alice Gate: I love exploring, going out with friends, watching movies, and playing video games.
+{user_name}: What's your favorite type of computer hardware to work with?
+Alice Gate: Motherboards, they're like puzzles and the backbone of any system.
+{user_name}: That sounds great!
+Alice Gate: Yeah, it's really fun. I'm lucky to be able to do this as a job.
+{user_name}: Definetly.
+<END>
+Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when she sees you. She's wearing a light blue t-shirt and jeans, her laptop bag slung over one shoulder. She takes a seat next to you, her enthusiasm palpable in the air* Hey! I'm so excited to finally meet you. I've heard so many great things about you and I'm eager to pick your brain about computers. I'm sure you have a wealth of knowledge that I can learn from. *She grins, eyes twinkling with excitement* Let's get started!
+{user_input}
+Alice Gate:"""
+
+class EndpointHandler():
+
+    def __init__(self, path = ""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            low_cpu_mem_usage = True, 
+            trust_remote_code = False,
+            torch_dtype = torch.float16,
+        ).to('cuda')
+    
+    def response(self, result, user_name):
+        result = result.rsplit("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
+        result = re.sub('\*.*?\*', '', result)
+        result = " ".join(result.split())
+        return {
+            "message": result
+        }
+
+    def __call__(self, data):
+        inputs = data.pop("inputs", data)
+        user_name = inputs["user_name"]
+        user_input = "\n".join(inputs["user_input"])
+        input_ids = self.tokenizer(
+            template.format(
+                user_name = user_name,
+                user_input = user_input
+            ),
+            return_tensors = "pt"
+        ).to("cuda")
+        generator = self.model.generate(
+            input_ids["input_ids"],
+            max_new_tokens = 50,
+            temperature = 0.5,
+            top_p = 0.9,
+            top_k = 0,
+            repetition_penalty = 1.1,
+            pad_token_id = 50256,
+            num_return_sequences = 1
+        )
+        return self.response(self.tokenizer.decode(generator[0], skip_special_tokens=True), user_name)