Update code/inference.py

code/inference.py

@@ -2,23 +2,6 @@ from transformers import AutoTokenizer
 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}"""
-
 def model_fn(model_dir):
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
     model = torch.load(f"{model_dir}/torch_model.pt")
@@ -26,14 +9,20 @@ def model_fn(model_dir):
 
 def predict_fn(input_data, load_list):
     model, tokenizer = load_list
-    inputs = input_data.pop("inputs", input_data)
-    user_name = inputs["user_name"]
-    user_input = "\n".join(inputs["user_input"])
-    prompt = template.format(
-        user_name = user_name,
-        user_input = user_input
-    )
-    input_ids = tokenizer(prompt + "\nAlice Gate:", return_tensors = "pt").to("cuda")
+    request_inputs = input_data.pop("inputs", input_data)
+    template = request_inputs["template"]
+    messages = request_inputs["messages"]
+    char_name = request_inputs["char_name"]
+    user_name = request_inputs["user_name"]
+    template = open(f"{template}.txt", "r").read()
+    user_input = [
+        "{name}: {message}".format(
+            name = char_name if (id["role"] == "AI") else user_name,
+            message = id["message"].strip()
+        ) for id in messages
+    ]
+    prompt = template.format(char_name = char_name, user_name = user_name, user_input = user_input)
+    input_ids = tokenizer(prompt + f"\n{char_name}:", return_tensors = "pt").to("cuda")
     encoded_output = model.generate(
         input_ids["input_ids"],
         max_new_tokens = 50,
@@ -45,13 +34,15 @@ def predict_fn(input_data, load_list):
         num_return_sequences = 1
     )
     decoded_output = tokenizer.decode(encoded_output[0], skip_special_tokens=True).replace(prompt,"")
-    decoded_output = decoded_output.split("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
+    decoded_output = decoded_output.split(f"{char_name}:", 1)[1].split(f"{user_name}:",1)[0].strip()
     parsed_result = re.sub('\*.*?\*', '', decoded_output).strip()
     if len(parsed_result) != 0: decoded_output = parsed_result
-    decoded_output = decoded_output.replace("*","")
-    decoded_output = " ".join(decoded_output.split())
+    decoded_output = " ".join(decoded_output.replace("*","").split())
     try:
         parsed_result = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1]
         if len(parsed_result) != 0: decoded_output = parsed_result
     except Exception: pass
-    return decoded_output
+    return {
+        "role": "AI",
+        "message": decoded_output
+    }