Update code/inference.py

code/inference.py

@@ -1,6 +1,7 @@
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from sagemaker_inference import content_types, decoder
 import torch
-import re
+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>
@@ -21,40 +22,30 @@ Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when
 Alice Gate:"""
 
 def model_fn(model_dir):
-  # Load model from HuggingFace Hub
-  tokenizer = AutoTokenizer.from_pretrained(model_dir)
-  model = torch.load(f"{model_dir}/torch_model.pt")
-  return model, tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(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 create_new_response(result, user_name):
-    result = result.rsplit("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
-    parsed_result = re.sub('\*.*?\*', '', result).strip()
-    result = parsed_result if len(parsed_result) != 0 else result.replace("*","")
-    result = " ".join(result.split())
-    try:
-        result = result[:[m.start() for m in re.finditer(r'[.!?]', result)][-1]+1]
-    except Exception: pass
-    return {
-        "message": result
-    }
+def output_fn(decoded_output, accept):
+    response_body = json.dumps({
+        "message": decoded_output
+    })
+    return response_body, accept
 
-def predict_fn(data, model_and_tokenizer):
-    # destruct model and tokenizer
-    model, tokenizer = model_and_tokenizer
-    
-    # Tokenize sentences
+def predict_fn(input_data, load_list):
+    model, tokenizer = load_list
+    inputs = 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,
-        return_tensors = "pt"
-    ).to("cuda")
-    generator = model.generate(
+    input_ids = tokenizer(prompt, return_tensors = "pt").to("cuda")
+    encoded_output = model.generate(
         input_ids["input_ids"],
         max_new_tokens = 50,
         temperature = 0.5,
@@ -64,4 +55,12 @@ def predict_fn(data, model_and_tokenizer):
         pad_token_id = 50256,
         num_return_sequences = 1
     )
-    return create_new_response(tokenizer.decode(generator[0], skip_special_tokens=True), user_name)
+    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()
+    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())
+    try:
+        decoded_output = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1]
+    except Exception: pass
+    return decoded_output