Update app.py

app.py

@@ -5,12 +5,14 @@ client = InferenceClient(
     "cognitivecomputations/dolphin-2.8-gemma-2b"
 )
 
+assistantStart = "<|im_start|>assistant\n"
+
 def format_prompt(system, message, history):
   prompt = f"<|im_start|>system\n{system}<|im_end|>\n"
   for user_prompt, bot_response in history:
     prompt += f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
     prompt += f"<|im_start|>assistant\n{bot_response}<|im_end|>\n"
-  prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+  prompt += f"<|im_start|>user\n{message}<|im_end|>\n{assistantStart}"
   return prompt
 
 def generate(
@@ -33,7 +35,9 @@ def generate(
 
     formatted_prompt = format_prompt(system_prompt, prompt, history)
     output = client.text_generation(formatted_prompt, **generate_kwargs)
-    return output
+    parsed = output.split(assistantStart, 1)
+    response = parsed[1] if len(parsed) > 1 else ""
+    return response
 
 
 additional_inputs=[