Improving UI

app.py

@@ -13,7 +13,6 @@ from scripts.UBAR_code.interaction import UBAR_interact
 from scripts.user_model_code.interaction import multiwoz_interact
 from scripts.UBAR_code.interaction.UBAR_interact import bcolors
 
-
 # Initialise agents
 UBAR_checkpoint_path = "epoch50_trloss0.59_gpt2"
 user_model_checkpoint_path = "MultiWOZ-full_checkpoint_step340k"
@@ -48,22 +47,41 @@ us_history = []
 self_play_history = []
 
 
+def reset_ds_state():
+    ds_history.clear()
+    sys_model.init_session()
+    return ds_history
+
+
+def reset_us_state():
+    us_history.clear()
+    user_model.init_session(ini_goal=current_goal)
+    return us_history
+
+
+def reset_self_play_state():
+    self_play_history.clear()
+    self_play_sys_model.init_session()
+    self_play_user_model.init_session(ini_goal=current_goal)
+    return self_play_history
+
+
 def change_goal():
     global curr_goal_idx
     curr_goal_idx = random.randint(0, n_goals - 1)
     current_goal = goals[curr_goal_idx]
-    user_model.init_session(ini_goal=current_goal)
+    us_history = reset_us_state()
     current_goal_yaml = yaml.dump(current_goal, default_flow_style=False)
-    return current_goal_yaml
+    return current_goal_yaml, us_history
 
 
 def change_sp_goal():
     global curr_sp_goal_idx
     curr_sp_goal_idx = random.randint(0, n_goals - 1)
     current_sp_goal = goals[curr_sp_goal_idx]
-    self_play_user_model.init_session(ini_goal=current_sp_goal)
+    self_play_history = reset_self_play_state()
     current_sp_goal_yaml = yaml.dump(current_sp_goal, default_flow_style=False)
-    return current_sp_goal_yaml
+    return current_sp_goal_yaml, self_play_history
 
 
 def ds_chatbot(user_utt):
@@ -87,10 +105,10 @@ def self_play():
     else:
         sys_response = self_play_history[-1][1]
 
-    user_utt = user_model.response(sys_response)
+    user_utt = self_play_user_model.response(sys_response)
 
     turn_id = len(self_play_history)
-    sys_response = sys_model.response(user_utt, turn_id)
+    sys_response = self_play_sys_model.response(user_utt, turn_id)
 
     self_play_history.append((user_utt, sys_response))
 
@@ -104,12 +122,20 @@ def self_play():
 block = gr.Blocks()
 
 with block:
-    gr.Markdown("# Demo User Simulator and Task-Oriented Dialogue System")
-    gr.Markdown("*Created by Alistair McLeay, with help from Professor Bill Byrne, Andy Tseng, and Alex Coca*")
+    gr.Markdown("#💬 Jointly Optimized Task-Oriented Dialogue System And User Simulator")
+    gr.Markdown(
+        "Created by [Alistair McLeay](https://alistairmcleay.com) for the [Masters in Machine Learning & Machine Intelligence at Cambridge University](https://www.mlmi.eng.cam.ac.uk/). <br/>\
+        Thank you to [Professor Bill Byrne](https://sites.google.com/view/bill-byrne/home) for his supervision and guidance. <br/> \
+        Thank you to [Andy Tseng](https://github.com/andy194673) and [Alex Coca](https://github.com/alexcoca) who provided code and guidance."
+    )
+    gr.Markdown(
+        "*Both Systems are trained on the [MultiWOZ dataset](https://github.com/budzianowski/multiwoz). <br/> \
+        Supported domains are 1. 🚆 Train, 2. 🏨 Hotel, 3. 🚕 Taxi, 4. 🚓 Police, 5. 🏣 Restaurant, 6. 🗿 Attraction, 7. 🏥 Hospital.*"
+    )
     with gr.Tabs():
         with gr.TabItem("Dialogue System"):
             gr.Markdown(
-                "This bot is a Task-Oriented Dialogue Systen. You are the user. Go ahead and try to book a train, or a hotel etc."
+                "This bot is a Task-Oriented Dialogue Systen. \nYou are the user. Go ahead and try to book a train, or a hotel etc."
             )
             with gr.Row():
                 ds_input_text = gr.inputs.Textbox(
@@ -117,34 +143,44 @@ with block:
                 )
                 ds_response = gr.outputs.Chatbot(label="Dialogue System Response")
             ds_button = gr.Button("Submit Message")
+            reset_ds_button = gr.Button("Reset Conversation")
 
         with gr.TabItem("User Simulator"):
             gr.Markdown(
-                "This bot is a User Simulator. You are the Task-Oriented Dialogue System. Your job is to help the user with their requests."
+                "This bot is a User Simulator. \nYou are the Task-Oriented Dialogue System. Your job is to help the user with their requests. \nIf you want the User Simulator to have a different goal press 'Generate New Goal'"
             )
-            new_goal_button = gr.Button("Generate Goal")
             with gr.Row():
                 us_input_text = gr.inputs.Textbox(
                     label="Dialogue System Message", placeholder="How can I help you today?"
                 )
                 us_response = gr.outputs.Chatbot(label="User Simulator Response")
-                current_goal_yaml = gr.outputs.Textbox(label="Current Goal (YAML)")
             us_button = gr.Button("Submit Message")
+            reset_us_button = gr.Button("Reset Conversation")
+            new_goal_button = gr.Button("Generate New Goal")
+            current_goal_yaml = gr.outputs.Textbox(label="New Goal (YAML)")
 
         with gr.TabItem("Self-Play"):
             gr.Markdown(
-                "In this scenario you define a goal and you then watch both agents interact where the User Simulator is trying to achieve the goal, and the Task-Oriented Dialogue System is trying to help the User Simulator do so."
+                "In this case both the User Simulator and the Task-Oriented Dialogue System are agents. \nGet them to interact by pressing 'Run Next Step' \nIf you want the User Simulator to have a different goal press 'Generate New Goal'"
             )
-            new_sp_goal_button = gr.Button("Generate Goal")
-            with gr.Row():
-                self_play_response = gr.outputs.Chatbot(label="Self-Play Output")
-                current_sp_goal_yaml = gr.outputs.Textbox(label="Current Goal (YAML)")
+            self_play_response = gr.outputs.Chatbot(label="Self-Play Output")
             self_play_button = gr.Button("Run Next Step")
+            reset_self_play_button = gr.Button("Reset Conversation")
+            new_sp_goal_button = gr.Button("Generate New Goal")
+            current_sp_goal_yaml = gr.outputs.Textbox(label="New Goal (YAML)")
+
+    gr.Markdown("## System Architecture Overview")
+    gr.Markdown(
+        "![System Architecture](https://huggingface.co/spaces/alistairmcleay/cambridge-masters-project/tree/main/system_architecture.png)"
+    )
 
     ds_button.click(ds_chatbot, ds_input_text, ds_response)
     us_button.click(us_chatbot, us_input_text, us_response)
     self_play_button.click(self_play, None, self_play_response)
-    new_goal_button.click(change_goal, None, current_goal_yaml)
-    new_sp_goal_button.click(change_sp_goal, None, current_sp_goal_yaml)
+    new_goal_button.click(change_goal, None, [current_goal_yaml, us_response])
+    new_sp_goal_button.click(change_sp_goal, None, [current_sp_goal_yaml, self_play_response])
+    reset_ds_button.click(reset_ds_state, None, ds_response)
+    reset_us_button.click(reset_us_state, None, us_response)
+    reset_self_play_button.click(reset_self_play_state, None, self_play_response)
 
 block.launch(share=True)