Update app.py

app.py

@@ -7,12 +7,12 @@ import gradio as gr
 from threading import Thread
 from huggingface_hub import login
 from icrawler.builtin import BingImageCrawler
-import os
 
 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 = os.environ.get("MODEL_ID")
 MODEL = "mistralai/Mistral-Nemo-Instruct-2407"
 
 TITLE = "<h1><center>Mistral-Nemo</center></h1>"
@@ -23,7 +23,6 @@ PLACEHOLDER = """
 </center>
 """
 
-
 CSS = """
 .duplicate-button {
     margin: auto !important;
@@ -45,118 +44,157 @@ h3 {
 footer{visibility: hidden}
 """
 
-device = "cuda" # for GPU usage or "cpu" for CPU usage
+device = "cuda"  # or "cpu"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
+# NOTE: fix_mistral_regex=True is recommended for this tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL, fix_mistral_regex=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,          # torch_dtype is deprecated in newer transformers
     device_map="auto",
-    ignore_mismatched_sizes=True)
+    ignore_mismatched_sizes=True
+)
+
+
+def _system_prompt_for(name: str) -> str:
+    return (
+        f"You should respond like {name}. "
+        "You should have a meaningful conversation. Don't repeat yourself. "
+        "You should only output your response. "
+        "You don't need to put quotes around what you're saying. "
+        "You don't need to put your name at the beginning of your response."
+    )
+
 
 @spaces.GPU()
 def get_response(conversation):
-    temperature  = 0.3
+    temperature = 0.3
     max_new_tokens = 512
     top_p = 1.0
-    top_k = 20 
+    top_k = 20
     penalty = 1.2
 
-    input_text=tokenizer.apply_chat_template(conversation, tokenize=False)
+    # 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(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-    
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        timeout=60.0,
+        skip_prompt=True,
+        skip_special_tokens=True,
+    )
+
     generate_kwargs = dict(
-        input_ids=inputs, 
-        max_new_tokens = max_new_tokens,
-        do_sample = False if temperature == 0 else True,
-        top_p = top_p,
-        top_k = top_k,
-        temperature = temperature,
+        input_ids=inputs,
+        max_new_tokens=max_new_tokens,
+        do_sample=False if temperature == 0 else True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
         streamer=streamer,
         repetition_penalty=penalty,
-        pad_token_id = 10,
+        pad_token_id=10,
     )
 
     with torch.no_grad():
         thread = Thread(target=model.generate, kwargs=generate_kwargs)
         thread.start()
-        
+
     buffer = ""
     for new_text in streamer:
         buffer += new_text
-        #yield buffer
 
     return buffer
-    
+
 
 @spaces.GPU()
 def stream_chat(history, character_a, character_b):
+    """
+    history: list of messages in messages format:
+      [{"role": "user" | "assistant", "content": "..."}, ...]
+
+    In the UI:
+      - 'user' messages are Character B
+      - 'assistant' messages are Character A
 
-    # From A to B
-    if(history==[]):
-        conversation_a = [{"role": "system", "content": "You should respond like " + character_b + ". You should have a meaningful conversation. Don't repeat yourself. You should only output your response. You don't need to put quotes around what your saying. You don't need to put your name at the beginning of your response."}]
-        for answer, prompt in history:
-            conversation_a.extend([
-                {"role": "user", "content": prompt}, 
-                {"role": "assistant", "content": answer},
-            ])
-        conversation_a.append({"role": "user", "content": "You are having a conversation with " +  character_a + ". Introduce yourself."})
-        response_b = get_response(conversation_a)
-        print('response_b', response_b)
+    Each button click:
+      1. B says something new (as 'user')
+      2. A replies (as 'assistant')
+    """
+    if history is None:
+        history = []
+
+    # ---------- B speaks ----------
+    if len(history) == 0:
+        # First turn: B introduces themselves to A
+        b_user_prompt = (
+            f"You are {character_b}. You are having a conversation with {character_a}. "
+            "Introduce yourself and start the conversation."
+        )
     else:
-        conversation_a = [{"role": "system", "content": "You should respond like " + character_b + ". You should have a meaningful conversation. Don't repeat yourself. You should only output your response. You don't need to put quotes around what your saying. You don't need to put your name at the beginning of your response."}]
-        for answer, prompt in history:
-            conversation_a.extend([
-                {"role": "user", "content": prompt}, 
-                {"role": "assistant", "content": answer},
-            ])
-        conversation_a.append({"role": "user", "content": history[-1][0] })
-        response_b = get_response(conversation_a)
-        print('response_b', response_b)
-        
-
-    # From B to A
-    conversation_b = [{"role": "system", "content": "You should respond like " + character_a + ". You should have a meaningful conversation. Don't repeat yourself. You should only output your response. You don't need to put quotes around what your saying. You don't need to put your name at the beginning of your response."}]
-    for prompt, answer in history:
-        conversation_b.extend([
-            {"role": "user", "content": prompt}, 
-            {"role": "assistant", "content": answer},
-        ])
-    conversation_b.append({"role": "user", "content": response_b })
-    response_a = get_response(conversation_b)
-    print('response_a', response_a)
-
-    # Append responses to history
-    history.append((response_b , response_a ))
-    print('history', history)
-    
-    return history
+        # Last A message is what B is responding to
+        last_msg = history[-1]
+        last_text = last_msg["content"]
+        b_user_prompt = (
+            f"{character_a} just said: \"{last_text}\". "
+            f"Respond in character as {character_b} and continue the conversation."
+        )
+
+    conv_for_b = [
+        {"role": "system", "content": _system_prompt_for(character_b)},
+        *history,
+        {"role": "user", "content": b_user_prompt},
+    ]
+    response_b = get_response(conv_for_b)
+    print("response_b:", response_b)
+
+    # ---------- A speaks ----------
+    conv_for_a = [
+        {"role": "system", "content": _system_prompt_for(character_a)},
+        *history,
+        {"role": "user", "content": response_b},
+    ]
+    response_a = get_response(conv_for_a)
+    print("response_a:", response_a)
+
+    # ---------- Append to chat history for UI ----------
+    new_history = history + [
+        {"role": "user", "content": response_b},       # B side
+        {"role": "assistant", "content": response_a},  # A side
+    ]
+    print("history:", new_history)
+
+    return new_history
 
 
 def get_img(keyword):
-    path = './' + keyword
+    path = "./" + keyword
     os.makedirs(path, exist_ok=True)
-    bing_crawler = BingImageCrawler(storage={'root_dir': path})
+    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', '.PNG', '.JPG', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff')):
+        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
 
+
 def set_characters(a, b):
     img_a = get_img(a)
-    img_b =  get_img(b)
-    return img_a, img_b, gr.update(avatar_images=(img_b, img_a))
-            
-chatbot = gr.Chatbot(height=600,show_label=False)
+    img_b = get_img(b)
+    # avatar_images=(user_avatar, assistant_avatar) => (B, A)
+    # Also reset chat history when characters change
+    return img_a, img_b, gr.update(avatar_images=(img_b, img_a), value=[])
+
+
+chatbot = gr.Chatbot(height=600, show_label=False)
 
-theme = gr.themes.Base(
-).set(
+theme = gr.themes.Base().set(
     body_background_fill="#e1fceb",
     color_accent_soft="#ffffff",
     border_color_accent="#e1fceb",
@@ -169,94 +207,47 @@ theme = gr.themes.Base(
 )
 
 with gr.Blocks() as demo:
-    gr.HTML("""
-
+    gr.HTML(
+        """
         <center> <h1> Bot vs Bot </h1> </center>
-    
         <center> by <a href="https://www.tonyassi.com/">Tony Assi</a> </center>
-        
         <center> <h3> Pick two icons and watch them have a conversation </h3> </center>
-    """)
+    """
+    )
+
     with gr.Row():
-        character_a = gr.Textbox(label='Character A', info='Choose a person', placeholder='Socrates, Edgar Allen Poe, George Washington')
-        character_b = gr.Textbox(label='Character B', info='Choose a person', placeholder='Madonna, Paris Hilton, Liza Minnelli')
-    character_button = gr.Button('Initiate Characters')
-    #characters_html = gr.HTML()
+        character_a = gr.Textbox(
+            label="Character A",
+            info="Choose a person",
+            placeholder="Socrates, Edgar Allen Poe, George Washington",
+        )
+        character_b = gr.Textbox(
+            label="Character B",
+            info="Choose a person",
+            placeholder="Madonna, Paris Hilton, Liza Minnelli",
+        )
+
+    character_button = gr.Button("Initiate Characters")
+
     with gr.Row():
         image_a = gr.Image(show_label=False, interactive=False)
-        gr.Markdown(' ')
+        gr.Markdown(" ")
         image_b = gr.Image(show_label=False, interactive=False)
-    
-
-    chat = gr.Chatbot(show_label=False)
-    submit_button = gr.Button()
-
-    character_button.click(set_characters, inputs=[character_a, character_b], outputs=[image_a, image_b, chat])
-    submit_button.click(stream_chat, inputs=[chat, character_a, character_b], outputs=[chat])
-    
-    
-    #gr.HTML(TITLE)
-    #gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
-    #video = gr.HTML("""
-    #  <video src="https://huggingface.co/spaces/ahjvdjf33/moods/resolve/main/Paris-Moods-2/passive3.mp4"
-    #         playsinline autoplay loop muted></video>
-    #""", elem_id="output_video")
-    """
-    gr.ChatInterface(
-        fn=stream_chat,
-        chatbot=chatbot,
-        textbox=gr.Textbox(visible=False),
-        fill_height=True,
-        retry_btn=None,
-        undo_btn=None,
-        clear_btn=None,
-        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False, visible=False),
-        additional_inputs=[
-            gr.Slider(
-                minimum=0,
-                maximum=1,
-                step=0.1,
-                value=0.3,
-                label="Temperature",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=128,
-                maximum=8192,
-                step=1,
-                value=1024,
-                label="Max new tokens",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=0.0,
-                maximum=1.0,
-                step=0.1,
-                value=1.0,
-                label="top_p",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=1,
-                maximum=20,
-                step=1,
-                value=20,
-                label="top_k",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=0.0,
-                maximum=2.0,
-                step=0.1,
-                value=1.2,
-                label="Repetition penalty",
-                render=False,
-            ),
-        ],
-    )
-    """
 
+    # IMPORTANT: type="messages" for Gradio 6
+    chat = gr.Chatbot(show_label=False, type="messages")
+    submit_button = gr.Button("Start Conversation")
 
+    character_button.click(
+        set_characters,
+        inputs=[character_a, character_b],
+        outputs=[image_a, image_b, chat],
+    )
+    submit_button.click(
+        stream_chat,
+        inputs=[chat, character_a, character_b],
+        outputs=[chat],
+    )
 
 if __name__ == "__main__":
     demo.launch(css=CSS, theme=theme)