added format promp and changed sotopia information accordion (#35)

app.py

@@ -12,7 +12,7 @@ from transformers import (
     BitsAndBytesConfig,
 )
 
-from utils import Agent, format_sotopia_prompt, get_starter_prompt
+from utils import Agent, format_sotopia_prompt, get_starter_prompt, format_bot_message
 from functools import cache
 
 DEPLOYED = os.getenv("DEPLOYED", "true").lower() == "true"
@@ -84,7 +84,7 @@ def introduction():
 
 
 def param_accordion(according_visible=True):
-    with gr.Accordion("Parameters", open=False, visible=according_visible):
+    with gr.Accordion("Parameters", open=True, visible=according_visible):
         model_name  = gr.Dropdown(
             choices=["cmu-lti/sotopia-pi-mistral-7b-BC_SR", "mistralai/Mistral-7B-Instruct-v0.1", "GPT3.5"],  # Example model choices
             value="cmu-lti/sotopia-pi-mistral-7b-BC_SR",  # Default value
@@ -116,45 +116,31 @@ def param_accordion(according_visible=True):
     return temperature, session_id, max_tokens, model_name 
 
 
-def sotopia_info_accordion(
-    human_agent, machine_agent, scenario, according_visible=True
-):
-    with gr.Accordion(
-        "Sotopia Information", open=False, visible=according_visible
-    ):
+def sotopia_info_accordion(human_agent, machine_agent, scenario, accordion_visible=True):
+    with gr.Accordion("Sotopia Information", open=accordion_visible):
         with gr.Row():
-            with gr.Column():
-                user_name = gr.Textbox(
-                    lines=1,
-                    label="username",
-                    value=human_agent.name,
-                    interactive=True,
-                    placeholder=f"{human_agent.name}: ",
-                    show_label=False,
-                    max_lines=1,
-                )
-            with gr.Column():
-                bot_name = gr.Textbox(
-                    lines=1,
-                    value=machine_agent.name,
-                    interactive=True,
-                    placeholder=f"{machine_agent.name}: ",
-                    show_label=False,
-                    max_lines=1,
-                    visible=False,
-                )
-            with gr.Column():
-                scenario = gr.Textbox(
-                    lines=4,
-                    value=scenario,
-                    interactive=False,
-                    placeholder="Scenario",
-                    show_label=False,
-                    max_lines=4,
-                    visible=False,
-                )
-    return user_name, bot_name, scenario
-
+            user_name = gr.Textbox(
+                lines=1,
+                label="Human Agent Name",
+                value=human_agent.name,
+                interactive=True,
+                placeholder="Enter human agent name",
+            )
+            bot_name = gr.Textbox(
+                lines=1,
+                label="Machine Agent Name",
+                value=machine_agent.name,
+                interactive=True,
+                placeholder="Enter machine agent name",
+            )
+            scenario_textbox = gr.Textbox(
+                lines=4,
+                label="Scenario Description",
+                value=scenario,
+                interactive=True,
+                placeholder="Enter scenario description",
+            )
+    return user_name, bot_name, scenario_textbox
 
 def instructions_accordion(instructions, according_visible=False):
     with gr.Accordion("Instructions", open=False, visible=according_visible):
@@ -206,7 +192,7 @@ def chat_tab():
         text_output = tokenizer.decode(
             output_tokens[0], skip_special_tokens=True
         )
-        return text_output
+        return format_bot_message(text_output)
 
     with gr.Column():
         with gr.Row():

utils.py

@@ -1,4 +1,16 @@
 from typing import List, Tuple
+import ast
+
+FORMAT_TEMPLATE = """ Your available action types are
+"none action speak non-verbal communication leave".
+Note: You can "leave" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, 3. you find it uninteresting/you lose your patience, 4. or for other reasons you want to leave.
+
+Please only generate a JSON string including the action type and the argument.
+Your action should follow the given format:
+\nAs an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}
+the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.
+\nHere is the output schema:\n```\n{\"description\": \"An interface for messages.\\nThere is only one required method: to_natural_language\", \"properties\": {\"action_type\": {\"title\": \"Action Type\", \"description\": \"whether to speak at this turn or choose to not do anything\", \"enum\": [\"none\", \"speak\", \"non-verbal communication\", \"action\", \"leave\"], \"type\": \"string\"}, \"argument\": {\"title\": \"Argument\", \"description\": \"the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action\", \"type\": \"string\"}}, \"required\": [\"action_type\", \"argument\"]}\n```\u001b[0m
+"""
 
 class Agent:
     def __init__(self, name, background, goal, secrets, personality):
@@ -36,6 +48,20 @@ def truncate_dialogue_history_to_length(dia_his, surpass_num, tokenizer):
     return trunc_dia
 
 
+def format_bot_message(bot_message) -> str:
+    json_response = ast.literal_eval(bot_message)
+    match json_response["action_type"]:
+        case "none":
+            return 'did nothing'
+        case "speak":
+            return json_response["argument"]
+        case "non-verbal communication":
+            return f'[{json_response["action_type"]}] {json_response["argument"]}'
+        case "action":
+            return f'[{json_response["action_type"]}] {json_response["argument"]}'
+        case "leave":
+            return 'left the conversation'
+    
 def dialogue_history_creation(history, user_name, bot_name):
     dialogue_history = ""
     for idx, turn in enumerate(history):
@@ -43,6 +69,8 @@ def dialogue_history_creation(history, user_name, bot_name):
         # TODOTODO (haofeiyu): we first assume that human talks first
         user_turn_idx = idx * 2
         bot_turn_idx = idx * 2 + 1
+        if not bot_message.startswith("["): # if action type == speak, need to add 'said: ' to be consistent with the dialog prompt
+            bot_message = "said :" + bot_message
         dialogue_history = f"{dialogue_history}\n\nTurn #{user_turn_idx}: {user_name}: {user_message}\n\nTurn #{bot_turn_idx}: {bot_name}: {bot_message}"
     last_turn_idx = len(history) * 2
     return dialogue_history, last_turn_idx
@@ -75,4 +103,4 @@ def format_sotopia_prompt(
     )
     prompt = f"{prompt}\n{dialogue_history}"
     prompt = f"{prompt}\n\nTurn #{last_turn_idx+1}: {user_name}: {message}\n.\nYou are at Turn #{last_turn_idx+2}."
-    return prompt
+    return prompt + FORMAT_TEMPLATE if use_format_guide else prompt