| import openai |
| import os |
| from langchain.embeddings.openai import OpenAIEmbeddings |
| from langchain.vectorstores import FAISS |
| import json |
| import pickle |
| from tqdm import tqdm |
|
|
| from synapse.envs.mind2web.env_utils import ( |
| load_json, |
| get_target_obs_and_act, |
| get_top_k_obs, |
| ) |
|
|
|
|
| def get_specifiers_from_sample(sample: dict) -> str: |
| website = sample["website"] |
| domain = sample["domain"] |
| subdomain = sample["subdomain"] |
| goal = sample["confirmed_task"] |
| specifier = ( |
| f"Website: {website}\nDomain: {domain}\nSubdomain: {subdomain}\nTask: {goal}" |
| ) |
|
|
| return specifier |
|
|
|
|
| def build_memory(memory_path: str, data_dir: str, top_k: int = 3): |
| openai.api_key = os.environ["OPENAI_API_KEY"] |
|
|
| score_path = "scores_all_data.pkl" |
| with open(os.path.join(data_dir, score_path), "rb") as f: |
| candidate_results = pickle.load(f) |
| candidate_scores = candidate_results["scores"] |
| candidate_ranks = candidate_results["ranks"] |
|
|
| specifiers = [] |
| exemplars = [] |
| samples = load_json(data_dir, "train") |
| for sample in tqdm(samples): |
| specifiers.append(get_specifiers_from_sample(sample)) |
| prev_obs = [] |
| prev_actions = [] |
| for s, act_repr in zip(sample["actions"], sample["action_reprs"]): |
| |
| sample_id = f"{sample['annotation_id']}_{s['action_uid']}" |
| for candidates in [s["pos_candidates"], s["neg_candidates"]]: |
| for candidate in candidates: |
| candidate_id = candidate["backend_node_id"] |
| candidate["score"] = candidate_scores[sample_id][candidate_id] |
| candidate["rank"] = candidate_ranks[sample_id][candidate_id] |
|
|
| _, target_act = get_target_obs_and_act(s) |
| target_obs, _ = get_top_k_obs(s, top_k) |
|
|
| if len(prev_obs) > 0: |
| prev_obs.append("Observation: `" + target_obs + "`") |
| else: |
| query = f"Task: {sample['confirmed_task']}\nTrajectory:\n" |
| prev_obs.append(query + "Observation: `" + target_obs + "`") |
| prev_actions.append("Action: `" + target_act + "` (" + act_repr + ")") |
|
|
| message = [] |
| for o, a in zip(prev_obs, prev_actions): |
| message.append({"role": "user", "content": o}) |
| message.append({"role": "assistant", "content": a}) |
| exemplars.append(message) |
|
|
| with open(os.path.join(memory_path, "exemplars.json"), "w") as f: |
| json.dump(exemplars, f, indent=2) |
|
|
| print(f"# of exemplars: {len(exemplars)}") |
|
|
| |
| embedding = OpenAIEmbeddings(model="text-embedding-ada-002") |
| metadatas = [{"name": i} for i in range(len(specifiers))] |
| memory = FAISS.from_texts( |
| texts=specifiers, |
| embedding=embedding, |
| metadatas=metadatas, |
| ) |
| memory.save_local(memory_path) |
|
|
|
|
| def retrieve_exemplar_name(memory, query: str, top_k) -> tuple[list[str], list[float]]: |
| docs_and_similarities = memory.similarity_search_with_score(query, top_k) |
| retrieved_exemplar_names = [] |
| scores = [] |
| for doc, score in docs_and_similarities: |
| retrieved_exemplar_names.append(doc.metadata["name"]) |
| scores.append(score) |
|
|
| return retrieved_exemplar_names, scores |
|
|
|
|
| def load_memory(memory_path): |
| embedding = OpenAIEmbeddings(model="text-embedding-ada-002") |
| memory = FAISS.load_local(memory_path, embedding) |
|
|
| return memory |
|
|
