Update code/inference.py

code/inference.py

@@ -17,8 +17,7 @@ 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:"""
+{user_input}"""
 
 def model_fn(model_dir):
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
@@ -34,7 +33,7 @@ def predict_fn(input_data, load_list):
         user_name = user_name,
         user_input = user_input
     )
-    input_ids = tokenizer(prompt, return_tensors = "pt").to("cuda")
+    input_ids = tokenizer(prompt + "\nAlice Gate:", return_tensors = "pt").to("cuda")
     encoded_output = model.generate(
         input_ids["input_ids"],
         max_new_tokens = 50,