fix and replace scenario and social goal agent name

sotopia-lab · Nov 17, 2023 · 2498cd9 · 2498cd9
1 parent 54f1e33
commit 2498cd9
Showing 1 changed file with 9 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
@@ -38,8 +38,11 @@ def get_agent_info(agent1_pk, agent2_pk, env_pk, agent_dict, env_dict):
     agent1_secret = agent_dict[agent1_pk]['secret']
     agent2_secret = agent_dict[agent2_pk]['secret']
 
-    agent1_goal = env_dict[env_pk]['agent_goals'][0].replace('<extra_info>', '')
-    agent2_goal = env_dict[env_pk]['agent_goals'][1].replace('<extra_info>', '')
+    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,
@@ -92,6 +95,10 @@ def fill_template(agent1_info, agent2_info, scenario):
         "Conversation Starts:\n.\nYou are at Turn #0."
     )
 
+    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(
         agent1_name=agent1_info['agent_name'],
         agent1_age=agent1_info['agent_age'],
@@ -115,8 +122,6 @@ def fill_template(agent1_info, agent2_info, scenario):
     return prompt + format_instruction
 
 
-
-
 with open('redis_json_data.json', 'r') as f:
     all_json_data = json.load(f)