Update app.py

app.py

@@ -13,7 +13,7 @@ bnb_config = BitsAndBytesConfig(
 # Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("Rorical/0-roleplay", return_dict=True, trust_remote_code=True, device_map="cuda", quantization_config=bnb_config)
 tokenizer = AutoTokenizer.from_pretrained("Rorical/0-roleplay", trust_remote_code=True)
-# tokenizer.chat_template = "{% for message in messages %}{{'<|im_start|>' + ((message['role'] + '\n') if message['role'] != '' else '') + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>星野\n' }}{% endif %}" # Be careful that this model used custom chat template.
+tokenizer.chat_template = "{% for message in messages %}{{'<|im_start|>' + ((message['role'] + ':\n') if message['role'] != '' else '') + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>星野:\n' }}{% endif %}" # Be careful that this model used custom chat template.
 
 # Define the response function
 @spaces.GPU
@@ -50,7 +50,7 @@ def respond(
     
     # Decode the generated response
     response = tokenizer.decode(generate_ids[0], skip_special_tokens=True)
-    response = response.split(f"老师\n{message}\n星野\n")[1]
+    response = response.split(f"老师:\n{message}\n星野:\n")[1]
     
     return response