Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
+import re  # for cleaning unwanted tags
 
 # Load model and tokenizer
 model_name = "GoofyLM/gonzalez-v1"
@@ -18,7 +19,6 @@ if tokenizer.pad_token is None:
 
 # Define a custom chat template if one is not available
 if tokenizer.chat_template is None:
-    # Basic ChatML-style template
     tokenizer.chat_template = "{% for message in messages %}\n{% if message['role'] == 'system' %}<|system|>\n{{ message['content'] }}\n{% elif message['role'] == 'user' %}<|user|>\n{{ message['content'] }}\n{% elif message['role'] == 'assistant' %}<|assistant|>\n{{ message['content'] }}\n{% endif %}\n{% endfor %}\n{% if add_generation_prompt %}<|assistant|>\n{% endif %}"
 
 def respond(
@@ -63,11 +63,12 @@ def respond(
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    # Stream response
+    # Stream response with unwanted patterns removed
     response = ""
     for token in streamer:
         response += token
-        yield response
+        cleaned_response = re.sub(r"<[|:].*?[|:]>", "", response)  # Remove <|msg|...> or <:...:...>
+        yield cleaned_response
 
 # Create Gradio interface
 demo = gr.ChatInterface(
@@ -81,4 +82,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()