Update app.py

app.py

@@ -67,8 +67,10 @@ def generate_response(prompt: str) -> dict:
     input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
     output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
     generated_response = (tokenizer.decode(output[0]))
+    # Extract the assistant's reply
+    assistant_reply = generated_response.split('\n')[1]
 
-    return {"user": prompt, "assistant": generated_response}
+    return {"user": prompt, "assistant": assistant_reply}
 
   
 def generate_prompt_response(persona_prompt: str, prompt: str) -> dict: