update

indiebot_arena/ui/battle.py

@@ -1,6 +1,7 @@
 import hashlib
 import os
 import random
+import re
 from collections.abc import Iterator
 from threading import Thread
 
@@ -19,6 +20,11 @@ base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
 docs_path = os.path.join(base_dir, "docs", "battle_header.md")
 
 
+def remove_chat_tokens(text: str) -> str:
+  pattern = re.compile(r'</?(?:start_of_turn|end_of_turn)>')
+  return pattern.sub('', text).strip()
+
+
 @spaces.GPU(duration=30)
 def generate(chat_history: list,
              model_id: str,
@@ -80,7 +86,8 @@ def bot1_response(history, model_id):
   history.append({"role": "assistant", "content": ""})
   conv_history = history[:-1]
   for text in generate(conv_history, model_id):
-    history[-1]["content"] = text
+    cleaned_text = remove_chat_tokens(text)
+    history[-1]["content"] = cleaned_text
     yield history, gr.update(interactive=True), gr.update(interactive=True)
 
 
@@ -89,7 +96,8 @@ def bot2_response(history, model_id):
   history.append({"role": "assistant", "content": ""})
   conv_history = history[:-1]
   for text in generate(conv_history, model_id):
-    history[-1]["content"] = text
+    cleaned_text = remove_chat_tokens(text)
+    history[-1]["content"] = cleaned_text
     yield history, gr.update(interactive=True), gr.update(interactive=True)
 
 

indiebot_arena/ui/playground.py

@@ -4,7 +4,7 @@ import gradio as gr
 
 from indiebot_arena.config import LOCAL_TESTING
 from indiebot_arena.service.arena_service import ArenaService
-from indiebot_arena.ui.battle import generate
+from indiebot_arena.ui.battle import generate, remove_chat_tokens
 from indiebot_arena.util.cache_manager import get_free_space_gb, clear_hf_cache
 
 DESCRIPTION = "### 💬 Playground"
@@ -29,7 +29,8 @@ def bot1_response(history, model_id):
   history.append({"role": "assistant", "content": ""})
   conv_history = history[:-1]
   for text in generate(conv_history, model_id):
-    history[-1]["content"] = text
+    cleaned_text = remove_chat_tokens(text)
+    history[-1]["content"] = cleaned_text
     yield history