Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
-import re  # For cleaning unwanted tokens
 
 # Load model and tokenizer
 model_name = "GoofyLM/gonzalez-v1"
@@ -19,17 +18,8 @@ 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 %}
-{% 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 %}"""
+    # 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(
     message,
@@ -69,30 +59,26 @@ def respond(
         pad_token_id=tokenizer.pad_token_id
     )
     
-    # Start generation in a separate thread
+    # Start generation in separate thread
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
-    # Stream response with filtering
+    
+    # Stream response
     response = ""
     for token in streamer:
         response += token
-        # 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()
+        yield response
 
 # Create Gradio interface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value="", label="System message"),
-        gr.Slider(1, 2048, value=72, label="Max new tokens"),
+        gr.Slider(1, 215, value=72, label="Max new tokens"),
         gr.Slider(0.1, 4.0, value=0.7, label="Temperature"),
         gr.Slider(0.1, 1.0, value=0.95, label="Top-p (nucleus sampling)"),
     ],
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo. launch()