Update code/inference.py

code/inference.py

@@ -1,7 +1,6 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from sagemaker_inference import content_types, decoder
+import re
 import torch
-import json
 
 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>
@@ -26,18 +25,9 @@ def model_fn(model_dir):
     model = torch.load(f"{model_dir}/torch_model.pt")
     return model, tokenizer
 
-def input_fn(self, input_data, content_type):
-    return decoder.decode(input_data, content_type)
-
-def output_fn(decoded_output, accept):
-    response_body = json.dumps({
-        "message": decoded_output
-    })
-    return response_body, accept
-
 def predict_fn(input_data, load_list):
     model, tokenizer = load_list
-    inputs = data.pop("inputs", input_data)
+    inputs = input_data.pop("inputs", input_data)
     user_name = inputs["user_name"]
     user_input = "\n".join(inputs["user_input"])
     prompt = template.format(
@@ -55,12 +45,14 @@ def predict_fn(input_data, load_list):
         pad_token_id = 50256,
         num_return_sequences = 1
     )
-    decoded_output = tokenizer.decode(encoded_output[0], skip_special_tokens=True)
-    decoded_output = result.rsplit("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
+    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()
     parsed_result = re.sub('\*.*?\*', '', decoded_output).strip()
-    decoded_output = parsed_result if len(parsed_result) != 0 else decoded_output.replace("*","")
-    decoded_output = " ".join(result.split())
+    if len(parsed_result) != 0: decoded_output = parsed_result
+    decoded_output = decoded_output.replace("*","")
+    decoded_output = " ".join(decoded_output.split())
     try:
-        decoded_output = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1]
+        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