Update app.py

app.py

@@ -12,7 +12,6 @@ MODEL_LIST = ["mistralai/Mistral-Nemo-Instruct-2407"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 login(token=HF_TOKEN)
 
-# MODEL = os.environ.get("MODEL_ID")
 MODEL = "mistralai/Mistral-Nemo-Instruct-2407"
 
 TITLE = "<h1><center>Mistral-Nemo</center></h1>"
@@ -46,13 +45,13 @@ footer{visibility: hidden}
 
 device = "cuda"  # or "cpu"
 
-# NOTE: fix_mistral_regex=True is recommended for this tokenizer
+# Recommended flag for this tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL, fix_mistral_regex=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     dtype=torch.bfloat16,          # torch_dtype is deprecated in newer transformers
     device_map="auto",
-    ignore_mismatched_sizes=True
+    ignore_mismatched_sizes=True,
 )
 
 
@@ -68,13 +67,15 @@ def _system_prompt_for(name: str) -> str:
 
 @spaces.GPU()
 def get_response(conversation):
+    """
+    conversation: list of {"role": "system" | "user" | "assistant", "content": str}
+    """
     temperature = 0.3
     max_new_tokens = 512
     top_p = 1.0
     top_k = 20
     penalty = 1.2
 
-    # conversation is already in messages format [{'role', 'content'}, ...]
     input_text = tokenizer.apply_chat_template(conversation, tokenize=False)
     inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(
@@ -110,21 +111,21 @@ def get_response(conversation):
 @spaces.GPU()
 def stream_chat(history, character_a, character_b):
     """
-    history: list of messages in messages format:
+    history: list of messages (messages format):
       [{"role": "user" | "assistant", "content": "..."}, ...]
 
     In the UI:
-      - 'user' messages are Character B
-      - 'assistant' messages are Character A
+      - user messages = Character B
+      - assistant messages = Character A
 
-    Each button click:
+    Each click:
       1. B says something new (as 'user')
       2. A replies (as 'assistant')
     """
     if history is None:
         history = []
 
-    # ---------- B speaks ----------
+    # ---------- B speaks (user side) ----------
     if len(history) == 0:
         # First turn: B introduces themselves to A
         b_user_prompt = (
@@ -132,7 +133,7 @@ def stream_chat(history, character_a, character_b):
             "Introduce yourself and start the conversation."
         )
     else:
-        # Last A message is what B is responding to
+        # Find last assistant message (A) to respond to
         last_msg = history[-1]
         last_text = last_msg["content"]
         b_user_prompt = (
@@ -148,7 +149,7 @@ def stream_chat(history, character_a, character_b):
     response_b = get_response(conv_for_b)
     print("response_b:", response_b)
 
-    # ---------- A speaks ----------
+    # ---------- A speaks (assistant side) ----------
     conv_for_a = [
         {"role": "system", "content": _system_prompt_for(character_a)},
         *history,
@@ -157,10 +158,10 @@ def stream_chat(history, character_a, character_b):
     response_a = get_response(conv_for_a)
     print("response_a:", response_a)
 
-    # ---------- Append to chat history for UI ----------
+    # ---------- Append to chat history ----------
     new_history = history + [
-        {"role": "user", "content": response_b},       # B side
-        {"role": "assistant", "content": response_a},  # A side
+        {"role": "user", "content": response_b},       # B's line
+        {"role": "assistant", "content": response_a},  # A's line
     ]
     print("history:", new_history)
 
@@ -173,14 +174,12 @@ def get_img(keyword):
     bing_crawler = BingImageCrawler(storage={"root_dir": path})
     bing_crawler.crawl(keyword=keyword, max_num=1)
 
-    # Look for image files in the folder
     for file_name in os.listdir(path):
         if file_name.lower().endswith(
             (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff")
         ):
             return os.path.join(path, file_name)
 
-    # If no image is found
     return None
 
 
@@ -188,7 +187,7 @@ def set_characters(a, b):
     img_a = get_img(a)
     img_b = get_img(b)
     # avatar_images=(user_avatar, assistant_avatar) => (B, A)
-    # Also reset chat history when characters change
+    # also reset chat history when characters change
     return img_a, img_b, gr.update(avatar_images=(img_b, img_a), value=[])
 
 
@@ -234,8 +233,8 @@ with gr.Blocks() as demo:
         gr.Markdown(" ")
         image_b = gr.Image(show_label=False, interactive=False)
 
-    # IMPORTANT: type="messages" for Gradio 6
-    chat = gr.Chatbot(show_label=False, type="messages")
+    # IMPORTANT: no 'type=' argument here; your Gradio build doesn't support it
+    chat = gr.Chatbot(show_label=False)
     submit_button = gr.Button("Start Conversation")
 
     character_button.click(