fix data processing agent1 and agent2 name bug

sotopia-lab · Nov 13, 2023 · 11a1a20 · 11a1a20
1 parent cee1227
commit 11a1a20
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/data_generate/step4_convert_json_to_gen_input.py b/data_generate/step4_convert_json_to_gen_input.py
@@ -41,6 +41,9 @@ def get_agent_info(agent1_pk, agent2_pk, env_pk, agent_dict, env_dict):
     agent1_goal = env_dict[env_pk]['agent_goals'][0].replace('<extra_info>', '').replace('</extra_info>', '')
     agent2_goal = env_dict[env_pk]['agent_goals'][1].replace('<extra_info>', '').replace('</extra_info>', '')
 
+    agent1_goal = agent1_goal.replace('Agent1', agent1_name).replace('Agent2', agent2_name).replace('agent2', agent2_name).replace('agent1', agent1_name)
+    agent2_goal = agent2_goal.replace('Agent2', agent2_name).replace('Agent1', agent1_name).replace('agent1', agent1_name).replace('agent2', agent2_name)
+
     agent1_info = {
         'agent_name': agent1_name,
         'agent_age': agent1_age,
@@ -92,10 +95,8 @@ def fill_template(agent1_info, agent2_info, scenario):
         "Conversation Starts:\n.\nYou are at Turn #0."
     )
 
-    scenario = scenario.replace('agent1', agent1_info['agent_name'])
-    scenario = scenario.replace('agent2', agent2_info['agent_name'])
-    scenario = scenario.replace('Agent1', agent1_info['agent_name'])
-    scenario = scenario.replace('Agent2', agent2_info['agent_name'])
+    scenario = scenario.replace('agent2', agent2_info['agent_name']).replace('agent1', agent1_info['agent_name'])
+    scenario = scenario.replace('Agent2', agent2_info['agent_name']).replace('Agent1', agent1_info['agent_name'])
 
 
     prompt = prompt_template.format(