Update code/inference.py

code/inference.py

@@ -3,25 +3,33 @@ import re
 import torch
 
 def model_fn(model_dir):
+
+    # Load Tokenizer, Model and Default template
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
     model = torch.load(f"{model_dir}/torch_model.pt")
-    return model, tokenizer
+    template = open(f"{model_dir}/default_template.txt","r").read()
+    return model, tokenizer, template
 
 def predict_fn(data, load_list):
-    model, tokenizer = load_list
+
+    # Get model, tokenzier and template from the model_fn
+    model, tokenizer, template = load_list
+
+    # Parse the input request into correct format to generate model input
     request_inputs = data.pop("inputs", 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 = "\n".join([
+    user_input = [
         "{name}: {message}".format(
             name = char_name if (id["role"] == "AI") else user_name,
             message = id["message"].strip()
         ) for id in messages
-    ])
+    ]
+    user_input = "\n".join([user_input])
     prompt = template.format(char_name = char_name, user_name = user_name, user_input = user_input)
+
+    # tokenize the model input, generate and decode output
     input_ids = tokenizer(prompt + f"\n{char_name}:", return_tensors = "pt").to("cuda")
     encoded_output = model.generate(
         input_ids["input_ids"],
@@ -34,6 +42,8 @@ def predict_fn(data, load_list):
         num_return_sequences = 1
     )
     decoded_output = tokenizer.decode(encoded_output[0], skip_special_tokens=True).replace(prompt,"")
+
+    # Parse the decoded output to the expected response
     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