pre-commit

sotopia-lab · Apr 7, 2024 · ca49e4a · ca49e4a
1 parent aa15c1b
commit ca49e4a
Showing 1 changed file with 7 additions and 19 deletions.
diff --git a/sotopia/database/serialization.py b/sotopia/database/serialization.py
@@ -84,9 +84,7 @@ def get_agent_name_to_social_goal_from_episode(
     episode: EpisodeLog,
 ) -> dict[str, str]:
     agents = [AgentProfile.get(agent) for agent in episode.agents]
-    agent_names = [
-        agent.first_name + " " + agent.last_name for agent in agents
-    ]
+    agent_names = [agent.first_name + " " + agent.last_name for agent in agents]
     environment = EnvironmentProfile.get(episode.environment)
     agent_goals = {
         agent_names[0]: environment.agent_goals[0],
@@ -129,22 +127,16 @@ def episodes_to_csv(
     """
     data = {
         "episode_id": [episode.pk for episode in episodes],
-        "scenario": [
-            get_scenario_from_episode(episode) for episode in episodes
-        ],
-        "codename": [
-            get_codename_from_episode(episode) for episode in episodes
-        ],
+        "scenario": [get_scenario_from_episode(episode) for episode in episodes],
+        "codename": [get_codename_from_episode(episode) for episode in episodes],
         "agents_background": [
             get_agents_background_from_episode(episode) for episode in episodes
         ],
         "social_goals": [
-            get_agent_name_to_social_goal_from_episode(episode)
-            for episode in episodes
+            get_agent_name_to_social_goal_from_episode(episode) for episode in episodes
         ],
         "social_interactions": [
-            get_social_interactions_from_episode(episode)
-            for episode in episodes
+            get_social_interactions_from_episode(episode) for episode in episodes
         ],
     }
     df = pd.DataFrame(data)
@@ -167,12 +159,8 @@ def episodes_to_json(
                 scenario=get_scenario_from_episode(episode),
                 codename=get_codename_from_episode(episode),
                 agents_background=get_agents_background_from_episode(episode),
-                social_goals=get_agent_name_to_social_goal_from_episode(
-                    episode
-                ),
-                social_interactions=get_social_interactions_from_episode(
-                    episode
-                ),
+                social_goals=get_agent_name_to_social_goal_from_episode(episode),
+                social_interactions=get_social_interactions_from_episode(episode),
                 reasoning=episode.reasoning,
                 rewards=get_rewards_from_episode(episode),
             )