Update visualization

README.md

@@ -2,7 +2,7 @@
 title: IDEFICS Playground
 emoji: 🐨
 colorFrom: red
-colorTo: indigo
+colorTo: green
 sdk: gradio
 sdk_version: 3.40.1
 app_file: app_dialogue.py

app_dialogue.py

@@ -21,7 +21,9 @@ MODELS = [  # TODO uncomment
 ]
 
 API_PATHS = {
-    "HuggingFaceM4/idefics-9b-instruct": "https://api-inference.huggingface.co/models/HuggingFaceM4/idefics-9b-instruct",
+    "HuggingFaceM4/idefics-9b-instruct": (
+        "https://api-inference.huggingface.co/models/HuggingFaceM4/idefics-9b-instruct"
+    ),
     "HuggingFaceM4/idefics-80b-instruct": (
         "https://api-inference.huggingface.co/models/HuggingFaceM4/idefics-80b-instruct"
     ),
@@ -411,7 +413,7 @@ with gr.Blocks(title="IDEFICS Playground", theme=gr.themes.Base()) as demo:
             repetition_penalty = gr.Slider(
                 minimum=0.0,
                 maximum=5.0,
-                value=1.1,
+                value=1.0,
                 step=0.01,
                 interactive=True,
                 label="Repetition penalty",
@@ -540,6 +542,11 @@ with gr.Blocks(title="IDEFICS Playground", theme=gr.themes.Base()) as demo:
             acc_text += text_token
             last_turn = chat_history.pop(-1)
             last_turn[-1] += acc_text
+            if last_turn[-1].endswith("\nUser"):
+                # Safeguard: sometimes (rarely), the model won't generate the token `<end_of_utterance>` and will go directly to generating `\nUser:`
+                # It will thus stop the generation on `\nUser:`. But when it exits, it will have already generated `\nUser`
+                # This post-processing ensures that we don't have an additional `\nUser` wandering around.
+                last_turn[-1] = last_turn[-1][:-5]
             chat_history.append(last_turn)
             yield "", None, chat_history
             acc_text = ""