import json import os from collections import defaultdict from copy import copy from functools import cache from typing import Dict, List, Literal, Tuple import jsonlines from tqdm import tqdm from sotopia_rl.prompter.sotopia_generate import generate_action from sotopia_rl.prompter.sotopia_utils import ( Agent, Environment, dialogue_history_prompt, get_context_prompt, ) ENVIRONMENT_PROFILES = "../../data/profiles/environmentprofiles_v1.jsonl" AGENT_PROFILES = "../../data/profiles/agentprofiles_v1.jsonl" RELATIONSHIP_PROFILES = "../../data/profiles/relationshipprofiles_v1.jsonl" Action = Literal[ "none", "action", "non-verbal communication", "speak", "leave" ] ACTION_TYPES: list[Action] = [ "none", "action", "non-verbal communication", "speak", "leave", ] TEMPERATURE = 0.7 @cache def get_sotopia_profiles( env_file: str = ENVIRONMENT_PROFILES, agent_file: str = AGENT_PROFILES, relationship_file: str = RELATIONSHIP_PROFILES, ) -> Tuple[ List[Tuple[str, str]], Dict[str, Environment], Dict[str, Agent], Dict[str, Dict[str, List[str]]], ]: with open(env_file, "r") as f: data = [json.loads(line) for line in f.readlines()] code_names_count: Dict[str, int] = defaultdict(int) environments = [] environment_dict = {} for profile in sorted(data, key=lambda x: x["codename"]): env_obj = Environment(profile) if profile["codename"] in code_names_count: environments.append( ( "{}_{:05d}".format( profile["codename"], code_names_count[profile["codename"]], ), env_obj._id, ) ) else: environments.append((profile["codename"], env_obj._id)) environment_dict[env_obj._id] = env_obj code_names_count[profile["codename"]] += 1 with open(agent_file, "r") as f: data = [json.loads(line) for line in f.readlines()] agent_dict = {} for profile in data: agent_obj = Agent(profile) agent_dict[agent_obj._id] = agent_obj with open(relationship_file, "r") as f: data = [json.loads(line) for line in f.readlines()] relationship_dict: Dict[str, Dict[str, List[str]]] = defaultdict( lambda: defaultdict(list) ) for profile in data: relationship_dict[profile["relationship"]][ profile["agent1_id"] ].append(profile["agent2_id"]) relationship_dict[profile["relationship"]][ profile["agent2_id"] ].append(profile["agent1_id"]) return environments, environment_dict, agent_dict, relationship_dict def run_chat( message: str, history: List[List[str]], bot_agent: Agent, user_agent: Agent, environment: Environment, model_selection: str, ) -> Tuple[str, str]: context = get_context_prompt(bot_agent, user_agent, environment) dialogue_history, next_turn_idx = dialogue_history_prompt( message, history, user_agent, bot_agent ) prompt_history = f"{context}{dialogue_history}" prompt, agent_action = generate_action( model_selection, prompt_history, next_turn_idx, ACTION_TYPES, bot_agent.name, TEMPERATURE, ) return prompt, agent_action.to_natural_language() def load_sotopia_pi_data( data_path: str, environment_dict: Dict[str, Environment], agent_dict: Dict[str, Agent], ) -> Tuple[List[Environment], List[Agent], List[Agent], List[str]]: with jsonlines.open(data_path, "r") as f: dataset = [line for line in f] envs = [] start_agents = [] end_agents = [] social_interactions = [] for data in tqdm(dataset): if ( "gpt" not in data["experiment_model_name_pairs"][1] or "gpt" not in data["experiment_model_name_pairs"][2] ): continue envs.append(environment_dict[data["environment_id"]]) start_agents.append(agent_dict[data["agent_ids"][0]]) end_agents.append(agent_dict[data["agent_ids"][1]]) social_interactions.append(data["social_interactions"]) return envs, start_agents, end_agents, social_interactions def generate_prompt_response_pairs( output_dir: str, model_selections: List[str], envs: List[Environment], start_agents: List[Agent], end_agents: List[Agent], social_interactions: List[str], num_episodes: int = 2, ) -> None: if not os.path.exists(output_dir): with open(output_dir, "w") as f: f.write("[]") with open(output_dir, "r") as f: result_pairs = json.load(f) all_ids = set() for result in result_pairs: all_ids.add( f"{result['environment_id']}_{result['start_agent_id']}_{result['end_agent_id']}" ) count = 0 for env, start_agent, end_agent, social_interaction in tqdm( zip(envs, start_agents, end_agents, social_interactions), total=len(envs), ): if f"{env._id}_{start_agent._id}_{end_agent._id}" in all_ids: count += 1 continue full_history = social_interaction.split("\n\n") curr_history = [] for i in range(0, len(full_history), 2): if i > 0: curr_history.append( [ f"Utterance {i//2 - 1} by " + full_history[i - 2], f"Utterance {i//2 - 1} by " + full_history[i - 1], ] ) message = f"Utterance {i//2} by " + full_history[i] try: prompt, response0 = run_chat( message, curr_history, end_agent, start_agent, env, model_selections[0], ) prompt, response1 = run_chat( message, curr_history, start_agent, end_agent, env, model_selections[1], ) except Exception: continue result_pairs.append( { "prompt": prompt, "message": message, "history": copy(curr_history), model_selections[0]: response0, model_selections[1]: response1, "environment_id": env._id, "start_agent_id": start_agent._id, "end_agent_id": end_agent._id, "scenario": env.scenario, "start_agent_name": start_agent.name, "end_agent_name": end_agent.name, "start_agent_goal": env.agent_goals[0], "end_agent_goal": env.agent_goals[1], } ) count += 1 all_ids.add(f"{env._id}_{start_agent._id}_{end_agent._id}") with open(output_dir, "w") as f: f.write(json.dumps(result_pairs, indent=4)) if count >= num_episodes: break if __name__ == "__main__": ( environments, environment_dict, agent_dict, relationship_dict, ) = get_sotopia_profiles() envs, start_agents, end_agents, social_interactions = load_sotopia_pi_data( "../../data/sotopia_pi_episodes.jsonl", environment_dict, agent_dict ) print("Loaded data with {} episodes".format(len(envs))) generate_prompt_response_pairs( "../../data/gpt35_gpt4_prompt_response_pairs.json", ["gpt-3.5-turbo", "gpt-4o"], envs, start_agents, end_agents, social_interactions, 100000, )