Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
-import re  # for cleaning unwanted tags
+import re  # For cleaning unwanted tokens
 
 # Load model and tokenizer
 model_name = "GoofyLM/gonzalez-v1"
@@ -19,7 +19,17 @@ if tokenizer.pad_token is None:
 
 # Define a custom chat template if one is not available
 if tokenizer.chat_template is None:
-    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 %}"
+    tokenizer.chat_template = """{% for message in messages %}
+{% if message['role'] == 'system' %}<|system|>
+{{ message['content'] }}
+{% elif message['role'] == 'user' %}<|user|>
+{{ message['content'] }}
+{% elif message['role'] == 'assistant' %}<|assistant|>
+{{ message['content'] }}
+{% endif %}
+{% endfor %}
+{% if add_generation_prompt %}<|assistant|>
+{% endif %}"""
 
 def respond(
     message,
@@ -59,16 +69,19 @@ def respond(
         pad_token_id=tokenizer.pad_token_id
     )
     
-    # Start generation in separate thread
+    # Start generation in a separate thread
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-    
-    # Stream response with unwanted patterns removed
+
+    # Stream response with filtering
     response = ""
     for token in streamer:
         response += token
-        cleaned_response = re.sub(r"<[|:].*?[|:]>", "", response)  # Remove <|msg|...> or <:...:...>
-        yield cleaned_response
+        # Remove angle-bracket tags (e.g. <|msg|...>, <username: ...>, <:...:...>)
+        cleaned_response = re.sub(r"<[^>]+>", "", response)
+        # Remove leading "Output:" if present (case-insensitive, line start)
+        cleaned_response = re.sub(r"(?i)^\s*output:\s*", "", cleaned_response)
+        yield cleaned_response.strip()
 
 # Create Gradio interface
 demo = gr.ChatInterface(