diff --git a/README.md b/README.md index 07b2e92a..ee4e743f 100644 --- a/README.md +++ b/README.md @@ -6,4 +6,5 @@ We split our overall framework into multiple parts 2. Together AI Finetuning --> Input the train and test data / Output model checkpoint 3. LLM Finetuning --> Input the train and test data / Output model checkpoint 4. LLM Deplyment --> Input LLM Finetuned model checkpoint / Output Deployable OpenAI type API -5. Eval --> Input model checkpoint / Output evaluation scores \ No newline at end of file +5. Eval --> Input model checkpoint / Output evaluation scores +6. Generate --> Input None / Output new data on redis \ No newline at end of file diff --git a/data_process/redis_data_filtering/prompt_reverse_engineering.py b/data_process/redis_data_filtering/prompt_reverse_engineering.py index aab51079..92eb9296 100644 --- a/data_process/redis_data_filtering/prompt_reverse_engineering.py +++ b/data_process/redis_data_filtering/prompt_reverse_engineering.py @@ -2,7 +2,7 @@ import os from collections import defaultdict from typing import Any, Dict, List, Tuple, Union, cast - +import transformers import pandas as pd import rich from rich.console import Console @@ -15,21 +15,14 @@ import enum #PROMPT_PREFIX = "Prompt after formatting:\n" - +MAX_TOKEN = 2048 PROMPT_TEMPLATE="""Prompt after formatting:\nImagine you are {agent}, your task is to act/speak as {agent} would, keeping in mind {agent}'s social goal. You can find {agent}'s background and goal in the 'Here is the context of the interaction' field. Note that {agent}'s secret and goal is only visible to you. You should try your best to achieve {agent}'s goal in a way that align with their character traits. Additionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before). {history}. -You are at Turn #{turn_number}. Your available action types are -{action_list}. -Note: You can "leave" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, 3. you find it uninteresting/you lose your patience, 4. or for other reasons you want to leave. - -Please only generate a JSON string including the action type and the argument. -Your action should follow the given format: -{format_instructions} -""" +You are at Turn #{turn_number}.""" #PYDANTIC_FORMAT_INSTRUCTIONS.format(schema=schema_str) FORMAT_TEMPLATE = """\nAs an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]} @@ -37,11 +30,37 @@ \nHere is the output schema:\n```\n{\"description\": \"An interface for messages.\\nThere is only one required method: to_natural_language\", \"properties\": {\"action_type\": {\"title\": \"Action Type\", \"description\": \"whether to speak at this turn or choose to not do anything\", \"enum\": [\"none\", \"speak\", \"non-verbal communication\", \"action\", \"leave\"], \"type\": \"string\"}, \"argument\": {\"title\": \"Argument\", \"description\": \"the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action\", \"type\": \"string\"}}, \"required\": [\"action_type\", \"argument\"]}\n```\u001b[0m""" +PROMPT_TEMPLATE_W_FORMAT="""Prompt after formatting:\nImagine you are {agent}, your task is to act/speak as {agent} would, keeping in mind {agent}'s social goal. +You can find {agent}'s background and goal in the 'Here is the context of the interaction' field. +Note that {agent}'s secret and goal is only visible to you. +You should try your best to achieve {agent}'s goal in a way that align with their character traits. +Additionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before). +{history}. +You are at Turn #{turn_number}. Your available action types are +"none action speak non-verbal communication leave". +Note: You can "leave" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, 3. you find it uninteresting/you lose your patience, 4. or for other reasons you want to leave. + +Please only generate a JSON string including the action type and the argument. +Your action should follow the given format: +\nAs an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]} +the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted. +\nHere is the output schema:\n```\n{\"description\": \"An interface for messages.\\nThere is only one required method: to_natural_language\", \"properties\": {\"action_type\": {\"title\": \"Action Type\", \"description\": \"whether to speak at this turn or choose to not do anything\", \"enum\": [\"none\", \"speak\", \"non-verbal communication\", \"action\", \"leave\"], \"type\": \"string\"}, \"argument\": {\"title\": \"Argument\", \"description\": \"the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action\", \"type\": \"string\"}}, \"required\": [\"action_type\", \"argument\"]}\n```\u001b[0m +""" # static ACTION_LIST = "none action speak non-verbal communication leave" #" ".join(ActionType) ACTION_REVERSE_MAP = {"left ": "leave", 'did n': 'none', 'said:': 'speak'} +MODEL_CHECKPOINT = "meta-llama/Llama-2-13b-chat-hf" +HF_TOKEN = "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" + + +TOKENIZER = transformers.AutoTokenizer.from_pretrained( + MODEL_CHECKPOINT, + padding = False, + truncation = False, + token=HF_TOKEN, + ) def to_natural_language(self) -> str: match self.action_type: @@ -101,10 +120,27 @@ def generate_result(msg): return str_result -def reverse_episode_log(epilog, later_speak=False): +def surpass_max_token_check(string, max_token=MAX_TOKEN, tokenizer=TOKENIZER): + prompt_tokens = len(tokenizer(string)['input_ids']) + return max(prompt_tokens - max_token, 0) + +def truncate_prompt_to_length(dia_his, surpass_num, tokenizer=TOKENIZER): + # context_len = len(tokenizer(context)['input_ids']) + dia_sen = dia_his.split("\n") + remove_len = 0 + i = 0 + while remove_len < surpass_num: + remove_len+=len(tokenizer(dia_sen[i])['input_ids']) + i+=1 + trunc_dia = "\n".join(p for p in dia_sen[i:]) + return trunc_dia + + +def reverse_episode_log(epilog, later_speak=False, include_format=False, max_token=MAX_TOKEN): episode_msg = epilog.messages # per episode agent_model = epilog.models[1] + promt_template = PROMPT_TEMPLATE_W_FORMAT if include_format else PROMPT_TEMPLATE if len(episode_msg) > 0: init_loop = episode_msg[0] @@ -131,7 +167,8 @@ def reverse_episode_log(epilog, later_speak=False): dial_history += "\n"+tpl[2] else: # for the first context, we don't need \n - dial_history += tpl[2] + context = tpl[2] + dial_history += context if tpl[0] == speaker: # if speaker is the agent, use what he said as result str_result = generate_result(tpl[2]) @@ -139,15 +176,22 @@ def reverse_episode_log(epilog, later_speak=False): if i%2 == turn_div: # take alternative turns as we always want to predict one agent, not both next_turn = i - prompt = PROMPT_TEMPLATE.format( - agent=speaker, history=dial_history, turn_number=next_turn, - action_list=ACTION_LIST, format_instructions=FORMAT_TEMPLATE) + prompt = promt_template.format( + agent=speaker, history=dial_history, turn_number=next_turn) + over_tokens = surpass_max_token_check(prompt) + if over_tokens > 0: + all_dial = dial_history[len(context):] + #print(all_dial) + trun_dial = truncate_prompt_to_length(all_dial, over_tokens) + prompt = promt_template.format( + agent=speaker, history=context+"\n"+trun_dial, turn_number=next_turn) turn_dic["prompt"] = prompt turn_dic['result'] = str_result prompt_result_instances.append(turn_dic) return prompt_result_instances + def parse_prompt_to_json(episode, dir, init_speak): prompt_result_instances = reverse_episode_log(episode, init_speak) diff --git a/data_process/redis_data_filtering/redis_filtering.py b/data_process/redis_data_filtering/redis_filtering.py index b4c95a74..690b7997 100644 --- a/data_process/redis_data_filtering/redis_filtering.py +++ b/data_process/redis_data_filtering/redis_filtering.py @@ -95,7 +95,7 @@ def goal_filter_per_env_agent(episodes): env_tpls.append((episodes[agent1_rank[i]], 0)) env_tpls.append((episodes[agent2_rank[i]], 1)) else: - if goal_score['agent1'][agent1_rank[i]] >= min(GOAL_KEEP_THRESHOD, agent1_avg) and (goal_score['agent2'][agent2_rank[i]] >= min(KEEP_THRESHOD, agent2_avg)): + if goal_score['agent1'][agent1_rank[i]] >= min(GOAL_KEEP_THRESHOD, agent1_avg) and (goal_score['agent2'][agent2_rank[i]] >= min(GOAL_KEEP_THRESHOD, agent2_avg)): env_tpls.append((episodes[agent1_rank[i]], 0)) env_tpls.append((episodes[agent1_rank[i]], 1)) diff --git a/llm_deploy/README.md b/llm_deploy/README.md index a6c72812..80e3f445 100644 --- a/llm_deploy/README.md +++ b/llm_deploy/README.md @@ -1,5 +1,173 @@ +## Deploy lora-finetuned model using vLLM variance + We need to use an unmerged branch to support deploying lora-finetuned model. (the forked repo is https://github.com/troph-team/vllm.git) Go to the vllm dir and pip install -e . -To notice https://github.com/vllm-project/vllm/issues/1283, need to modify the config file to "== 2.0.1" and the pytorch version if facing with CUDA version error. \ No newline at end of file +To notice https://github.com/vllm-project/vllm/issues/1283, need to modify the config file to "== 2.0.1" and the pytorch version if facing with CUDA version error. + + +## Setting up Babel server +### Login with SSH key +Add public ed25519 key to server +```bash +ssh-copy-id -i ~/.ssh/id_ed25519.pub @ +``` +Config SSH file +```bash + Host + HostName + User + IdentityFile ~/.ssh/id_ed25519 +``` +Login babel with SSH key +```bash +ssh +``` + +### Connecting to a compute node +Jump from login node to compute node +```bash +srun --pty bash +``` +Check if you can access the /data/folder +```bash +cd /data/datasets/ +``` + +### Config environment on the compute node +Install miniconda +```bash +wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh +bash Miniconda3-latest-Linux-x86_64.sh +conda init +conda create --name myenv +conda activate myenv +# conda deactivate +``` +Install vllm packages +```bash +conda install pip +pip install vllm +``` +Install fastchat packages +```bash +conda install pip +git clone https://github.com/lm-sys/FastChat.git +cd FastChat +pip3 install --upgrade pip +pip3 install "fschat[model_worker,webui]" +``` +Submit gpu request and open a an interactive terminal +```bash +srun --gres=gpu:1 --time=1-00:00:00 --mem=80G --pty $SHELL +conda activate myenv +``` +Some useful commands for checking gpu jobs +```bash +# check slurm status +squeue -l +# check gpu status +nvidia-smi +# check gpu usage +pip install gpustat +watch -n 1 gpustat +# quit slurm jobs +scancel job_id +# connect to compute node directly +ssh -J babel babel-x-xx +``` + +### Install cuda-toolkit (optional) +Due to the issue with vllm: https://github.com/vllm-project/vllm/issues/1283, we need to use cuda-toolkit=11.7.0 that is compatible with Pytorch 2.0.1. +Install cuda-toolkit=11.7.0 on conda environment +```bash +conda install -c "nvidia/label/cuda-11.7.0" cuda-toolkit +``` +Check cuda-toolkit version +```bash +nvcc -V +``` + +## Deploy models on Babel via FastChat API server +Implement the following python commands in three separate interactive terminal windows: +```bash +python3 -m fastchat.serve.controller +python3 -m fastchat.serve.model_worker --model-path model-checkpoint +python3 -m fastchat.serve.openai_api_server --host localhost --port 8000 +``` +Call model checkpoint API +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "model-checkpoint", + "prompt": "San Francisco is a", + "max_tokens": 7, + "temperature": 0 + }' +``` +*Sample output:* +```JSON +{"id":"cmpl-GGvKBiZFdFLzPq2HdtuxbC","object":"text_completion","created":1698692212,"model":"checkpoint-4525","choices":[{"index":0,"text":"city that is known for its icon","logprobs":null,"finish_reason":"length"}],"usage":{"prompt_tokens":5,"total_tokens":11,"completion_tokens":6}} +``` + +## Deploy models on Babel via vllm API server +Start vLLM surver with model checkpoint +```bash +python -m vllm.entrypoints.openai.api_server --model model_checkpoint/ +``` +Call model checkpoint API +```bash +curl http://localhost:8000/v1/models +``` +*Sample output:* +```JSON +{"object":"list","data":[{"id":"Mistral-7B-Instruct-v0.1/","object":"model","created":1697599903,"owned_by":"vllm","root":"Mistral-7B-Instruct-v0.1/","parent":null,"permission":[{"id":"modelperm-d415ecf6362a4f818090eb6428e0cac9","object":"model_permission","created":1697599903,"allow_create_engine":false,"allow_sampling":true,"allow_logprobs":true,"allow_search_indices":false,"allow_view":true,"allow_fine_tuning":false,"organization":"*","group":null,"is_blocking":false}]}]} +``` +Inference model checkpoint API +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "model_checkpoint", + "prompt": "San Francisco is a", + "max_tokens": 7, + "temperature": 0 + }' +``` +*Sample output:* +```JSON +{"id":"cmpl-bf7552957a8a4bd89186051c40c52de4","object":"text_completion","created":3600699,"model":"Mistral-7B-Instruct-v0.1/","choices":[{"index":0,"text":" city that is known for its icon","logprobs":null,"finish_reason":"length"}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}} +``` + +## Access deployed Babel server on a local machine +Construct ssh tunnel between babel login node and babel compute node with hosted model +```bash +ssh -N -L 7662:localhost:8000 username@babel-x-xx +``` +The above command creates a localhost:7662 server on bable login node which connects to localhost:8000 on compute node. + +Construct ssh tunnel between local machine and babel login node +```bash +ssh -N -L 8001:localhost:7662 username@ +``` +The above command creates a localhost:8001 server on your local machine which connects to localhost:7662 on babel login node. + +Call hosted model on local machine +```bash +curl http://localhost:8001/v1/models +``` +If the above command runs successfully, you should be able to use REST API on your local machine. + +(optional) If you fail in building the ssh tunnel, you may add `-v` to the ssh command to see what went wrong. + + + + +## Userful resource links for babel +1. https://hpc.lti.cs.cmu.edu/wiki/index.php?title=BABEL#Cluster_Architecture +2. https://hpc.lti.cs.cmu.edu/wiki/index.php?title=VSCode +3. https://hpc.lti.cs.cmu.edu/wiki/index.php?title=Training_Material +4. https://hpc.lti.cs.cmu.edu/wiki/index.php?title=Connecting_to_the_Cluster#Copying_Data_to_Compute_Nodes + diff --git a/llm_generate/README.md b/llm_generate/README.md new file mode 100644 index 00000000..869e5537 --- /dev/null +++ b/llm_generate/README.md @@ -0,0 +1,9 @@ +# Data Generation + +For the first step, we generate envProfile (including scenario / social goal / relationship restriction) based on inspiring prompt. + +For the second step, we put the original agentProfile and relationshipProfile into our new redis database + +For the third step, we combine them together to be combos based on conditiona sampling (the restriction is the relationship) + +All the EnvProfile (new generated), AgentProfile (sotopia original), RelationshipProfile (sotopia original), and envagentcombo are on the redis database that is new created. \ No newline at end of file diff --git a/llm_generate/generate_specific_envs.py b/llm_generate/generate_specific_envs.py new file mode 100644 index 00000000..94ab8014 --- /dev/null +++ b/llm_generate/generate_specific_envs.py @@ -0,0 +1,135 @@ +"""This file is used to generate specific environments based on existing +datasets. The generation functions below should call agenerate_env_profile +in `sotopia/generation_utils/generate.py` with the appropriate parameters. +Here are the datasets we have so far: +1. Mutual-Friend (https://huggingface.co/datasets/mutual_friends) +""" +import asyncio +from typing import Hashable + +import datasets +import names +import numpy as np +from datasets import DatasetDict, load_dataset + +from generate import ( + ListOfStrOutputParser, + StrOutputParser, + agenerate, + generate, +) + + +async def generate_mutual_friend_envs() -> tuple[str, list[str]]: + """Generate environments based on the mutual-friend dataset.""" + mutual_friend_dataset: DatasetDict = load_dataset("mutual_friends") + all_data = mutual_friend_dataset["train"] + # sample one datum from all data + datum = np.random.choice(all_data) + friends = datum["scenario_kbs"] + num_of_friends_in_total = sum(map(len, friends)) + # generate names for the friends + set_of_names = set() + for _ in range(num_of_friends_in_total): + name = names.get_first_name() + while name in set_of_names: + name = names.get_first_name() + set_of_names.add(name) + list_of_names = list(set_of_names) + friend_map: dict[tuple[str, ...], str] = {} + friend_list_map: list[list[str]] = [[] for _ in range(len(friends))] + friend_description_keys: list[str] = datum["scenario_attributes"]["name"] + name_pointer = 0 + for i, friends_array in enumerate(friends): + for friend in friends_array: + assert ( + len(friend) == 2 + ) # in [[key1, key2, ...], [value1, value2, ...]] format + if not tuple(friend[1]) in friend_map: + friend_map[tuple(friend[1])] = list_of_names[name_pointer] + name_pointer += 1 + friend_list_map[i].append(friend_map[tuple(friend[1])]) + friend_set_map: list[set[str]] = [ + set(friend_list) for friend_list in friend_list_map + ] + common_friends = [] + for friend_description, friend_name in friend_map.items(): + if all([friend_name in friend_set for friend_set in friend_set_map]): + common_friends.append(friend_name) + scenario = ( + f'{len(friends)} strangers are meeting at a party.

They have {len(common_friends)} common friends: ' + f"{', '.join(common_friends[:-1])}" + + (" and " if len(common_friends) > 1 else "") + + common_friends[-1] + + ".

" + ) + goals: list[str] = [] + for friends_array in friends: + template = f"You are trying to figure out whether you have a mutual friend with the other person. \n" + template += f" You know the following friends" + for friend in friends_array: + friend_name = friend_map[tuple(friend[1])] + friend_description = friend[1] + template += f" {friend_name}: {' '.join([(i + ': ' + j + ' ') if i != 'Name' else '' for i, j in zip(friend[0], friend_description)])}\n" + template += f"" + goals.append(template) + + return scenario, goals + + +async def generate_craigslist_bargains_envs() -> tuple[str, list[str]]: + """Generate environments based on the craigslist_bargains dataset.""" + craigslist_bargains_dataset: DatasetDict = load_dataset( + "craigslist_bargains" + ) + all_data = craigslist_bargains_dataset["train"] + # sample one datum from all data + datum = np.random.choice(all_data) + scenario = generate( + model_name="gpt-4", + template="The following sentence is automatically generated with the following" + 'template: "One person is selling for , another person is' + 'trying to buy it. Here is the description of the item: ." with item = {title}, ' + "price={price}, and description={description} Please make the sentence" + "fluent and natural.", + input_values={ + "title": datum["items"]["Title"][0], + "price": datum["items"]["Price"][0], + "description": datum["items"]["Description"][0], + }, + output_parser=StrOutputParser(), + ) + + goals: list[str] = [] + for i in range(2): + if datum["agent_info"]["Role"][i] == "seller": + markup_ratio = np.random.exponential(0.5) + datum["agent_info"]["Target"][i] = datum["items"]["Price"][0] / ( + 1 + markup_ratio + ) + goal = generate( + model_name="gpt-4", + template="The following sentence is automatically generated with the following" + 'template: "You want to this item. Your target price ' + "is $ (round up to two decimals). You will get penalty if you sell or buy it " + "for a price that is significantly lower than (if is seller) or significantly" + "higher than (if is buyer) the target price, but will get bonus if you successfully " + "sell it higher than the target price (if is seller) or buy it for lower than" + 'the target price (if is buyer)." ' + "with role = {role} and price = {price}. Please make the sentence" + "fluent and natural. Do not change the original meaning of the sentence.", + input_values={ + "role": datum["agent_info"]["Role"][i], + "price": datum["agent_info"]["Target"][i], + }, + output_parser=StrOutputParser(), + ) + goals.append(goal) + + return scenario, goals + + +if __name__ == '__main__': + for i in range(10): + scenario, goals = asyncio.run(generate_mutual_friend_envs()) + import pdb; pdb.set_trace() \ No newline at end of file diff --git a/llm_generate/requirments.txt b/llm_generate/requirments.txt new file mode 100644 index 00000000..804f2e66 --- /dev/null +++ b/llm_generate/requirments.txt @@ -0,0 +1 @@ +sotopia \ No newline at end of file diff --git a/llm_generate/step1_generate_env_profile.py b/llm_generate/step1_generate_env_profile.py new file mode 100644 index 00000000..4e9f9824 --- /dev/null +++ b/llm_generate/step1_generate_env_profile.py @@ -0,0 +1,51 @@ +import asyncio +import random +from typing import TypeVar +from tqdm import tqdm + +import pandas as pd +import rich +from pydantic import BaseModel + +from sotopia.database import EnvironmentProfile +from sotopia.generation_utils.generate import agenerate_env_profile + +random.seed(41) + +env_borrowMoney = EnvironmentProfile.find( + EnvironmentProfile.codename == "borrow_money" +).all()[0] +env_roadtrip = EnvironmentProfile.find( + EnvironmentProfile.codename == "take_turns" +).all()[0] +env_prisonerDillema = EnvironmentProfile.find( + EnvironmentProfile.codename == "prison_dilemma" +).all()[0] + +examples = f"{env_borrowMoney.json()}\n\n{env_roadtrip.json()}\n\n{env_prisonerDillema.json()}" + +ins_prompts = pd.read_csv("./inspirational_prompt_for_env.csv") +prompts = ins_prompts["prompt"].tolist() + +T = TypeVar("T", bound=BaseModel) + + +def pydantics_to_csv(filename: str, data: list[T]) -> None: + pd.DataFrame([item.dict() for item in data]).to_csv(filename, index=False) + + +backgrounds = [] +for prompt in tqdm(prompts): + rich.print(prompt) + background, prompt_full = asyncio.run( + agenerate_env_profile( + model_name="gpt-4", + inspiration_prompt=prompt, + examples=examples, + ) + ) + rich.print(background) + rich.print(prompt_full) + backgrounds.append(background) + + pydantics_to_csv("./backgrounds.csv", backgrounds) \ No newline at end of file diff --git a/llm_generate/step2_push_agent_relationship_env_to_db.py b/llm_generate/step2_push_agent_relationship_env_to_db.py new file mode 100644 index 00000000..7a8fe8e8 --- /dev/null +++ b/llm_generate/step2_push_agent_relationship_env_to_db.py @@ -0,0 +1,150 @@ +import ast +import sys +from typing import Any, cast + +import pandas as pd +from redis_om import Migrator + +from sotopia.database.persistent_profile import ( + AgentProfile, + EnvironmentProfile, + RelationshipProfile, +) +from sotopia.database.env_agent_combo_storage import EnvAgentComboStorage +from sotopia.samplers import ConstraintBasedSampler +from sotopia.messages import AgentAction, Observation +from sotopia.agents import LLMAgent + + + +def add_agent_to_database(**kwargs: dict[str, Any]) -> None: + agent = AgentProfile(**kwargs) + agent.save() + + +def add_agents_to_database(agents: list[dict[str, Any]]) -> None: + for agent in agents: + add_agent_to_database(**agent) + + +def retrieve_agent_by_first_name(first_name: str) -> AgentProfile: + result = AgentProfile.find(AgentProfile.first_name == first_name).all() + if len(result) == 0: + raise ValueError(f"Agent with first name {first_name} not found") + elif len(result) > 1: + raise ValueError(f"Multiple agents with first name {first_name} found") + else: + assert isinstance(result[0], AgentProfile) + return result[0] + + +def add_env_profile(**kwargs: dict[str, Any]) -> None: + env_profile = EnvironmentProfile(**kwargs) + env_profile.save() + + +def add_env_profiles(env_profiles: list[dict[str, Any]]) -> None: + for env_profile in env_profiles: + add_env_profile(**env_profile) + + +def add_relationship_profile(**kwargs: dict[str, Any]) -> None: + relationship_profile = RelationshipProfile(**kwargs) + relationship_profile.save() + + +def add_relationship_profiles( + relationship_profiles: list[dict[str, Any]] +) -> None: + for relationship_profile in relationship_profiles: + add_relationship_profile(**relationship_profile) + + +def delete_all_agents() -> None: + pks = AgentProfile.all_pks() + pks_list = list(pks) + for id in pks: + AgentProfile.delete(id) + + +def delete_all_env_profiles() -> None: + pks = EnvironmentProfile.all_pks() + #for id in pks: + # EnvironmentProfile.delete(id) + + +def delete_all_relationships() -> None: + pks = list(RelationshipProfile.all_pks()) + #for id in pks: + # RelationshipProfile.delete(id) + pks = list(RelationshipProfile.all_pks()) + print("Relationships deleted, all relationships: ", len(list(pks))) + + +def sample_env_agent_combo_and_push_to_db(env_id: str) -> None: + sampler = ConstraintBasedSampler[Observation, AgentAction]( + env_candidates=[env_id] + ) + try: + env_agent_combo_list = list( + sampler.sample(agent_classes=[LLMAgent] * 2, replacement=False) + ) + except: + return + print(len(env_agent_combo_list)) + for env, agent in env_agent_combo_list: + EnvAgentComboStorage( + env_id=env.profile.pk, + agent_ids=[agent[0].profile.pk, agent[1].profile.pk], + ).save() + + +def relationship_map(relationship: str) -> int: + return int(eval(relationship)) + + +if __name__ == "__main__": + assert ( + len(sys.argv) == 3 + ), "Please provide a csv file with agent or environment profiles, and the type of profile (agent or environment)" + df = pd.read_csv(sys.argv[1]) + type = sys.argv[2] + if type == "agent": + agents = cast(list[dict[str, Any]], df.to_dict(orient="records")) + for agent in agents: + agent["age"] = int(agent["age"]) + agent["moral_values"] = agent["moral_values"].split(",") + agent["schwartz_personal_values"] = agent[ + "schwartz_personal_values" + ].split(",") + add_agents_to_database(agents) + Migrator().run() + elif type == "environment": + df = df[ + [ + "codename", + "scenario", + "agent_goals", + "relationship", + "age_constraint", + "occupation_constraint", + "source", + ] + ] + envs = cast(list[dict[str, Any]], df.to_dict(orient="records")) + for env in envs: + env["agent_goals"] = ast.literal_eval(env["agent_goals"]) + assert isinstance(env["relationship"], int) + Migrator().run() + elif type == "relationship": + relationships = cast( + list[dict[str, Any]], df.to_dict(orient="records") + ) + for relationship in relationships: + assert isinstance(relationship["relationship"], int) + add_relationship_profiles(relationships) + Migrator().run() + elif type == 'agentenvcombo': + env_ids = list(EnvironmentProfile.all_pks()) + for env_id in env_ids: + sample_env_agent_combo_and_push_to_db(env_id) \ No newline at end of file diff --git a/llm_generate/test_redis1.py b/llm_generate/test_redis1.py new file mode 100644 index 00000000..34c18a52 --- /dev/null +++ b/llm_generate/test_redis1.py @@ -0,0 +1,10 @@ +import redis + +r = redis.Redis( + host='us1-normal-burro-37804.upstash.io', + port=37804, + password='a870a438f928424bb507d5895b3ab3fc' +) + +r.set('foo', 'bar') +print(r.get('foo')) \ No newline at end of file diff --git a/llm_generate/test_redis2.py b/llm_generate/test_redis2.py new file mode 100644 index 00000000..06c691c9 --- /dev/null +++ b/llm_generate/test_redis2.py @@ -0,0 +1,11 @@ +from redis_om import JsonModel, get_redis_connection + +class Person(JsonModel): + name: str + age: int + +# Create an instance of your model +person = Person(name="John", age=30) + +# Save to Redis with a specific key +person.save() \ No newline at end of file