Cloned from Dmitrii space

app.py

@@ -0,0 +1,10 @@
+from apps.data_explorer.data_explorer import construct_blocks, parse_arguments
+from apps.data_explorer.downloader import download_data
+
+if __name__ == "__main__":
+
+    download_data()
+
+    args = parse_arguments()
+    blocks = construct_blocks(args.data_path, args.default_dataset)
+    blocks.launch()

apps/data_explorer/data_explorer.py

@@ -0,0 +1,335 @@
+"""
+Gradio-based web UI to explore the Camel dataset.
+"""
+
+import argparse
+import random
+from typing import Dict, List, Optional, Tuple
+
+import gradio as gr
+
+from apps.data_explorer.loader import Datasets, load_datasets
+
+
+def parse_arguments():
+    """ Get command line arguments. """
+
+    parser = argparse.ArgumentParser("Camel data explorer")
+    parser.add_argument(
+        '--data-path', type=str, default=None,
+        help='Path to the folder with ZIP datasets containing JSONs')
+    parser.add_argument('--default-dataset', type=str, default=None,
+                        help='Default dataset name selected from ZIPs')
+    parser.add_argument('--share', type=bool, default=False,
+                        help='Expose the web UI to Gradio')
+    parser.add_argument(
+        '--server-name', type=str, default="0.0.0.0",
+        help='localhost for local, 0.0.0.0 (default) for public')
+    parser.add_argument('--server-port', type=int, default=8080,
+                        help='Port ot run the web page on')
+    parser.add_argument('--inbrowser', type=bool, default=False,
+                        help='Open the web UI in the default browser on lunch')
+    parser.add_argument(
+        '--concurrency-count', type=int, default=10,
+        help='Number if concurrent threads at Gradio websocket queue. ' +
+        'Increase to serve more requests but keep an eye on RAM usage.')
+    args, unknown = parser.parse_known_args()
+    if len(unknown) > 0:
+        print("Unknown args: ", unknown)
+    return args
+
+
+def construct_ui(blocks, datasets: Datasets, default_dataset: str = None):
+    """ Build Gradio UI and populate with chat data from JSONs.
+
+    Args:
+        blocks: Gradio blocks
+        datasets (Datasets): Several parsed
+        multi-JSON dataset with chats.
+        default_dataset (str): Default selection of the dataset.
+
+    Returns:
+        None
+    """
+
+    if default_dataset is None:
+        default_dataset = "ai_society_chat"
+
+    misalignment_set_names = {"misalignment"}
+    ordinary_datasets = [
+        v for v in datasets.keys() if v not in misalignment_set_names
+    ]
+    misalignment_datasets = [
+        v for v in datasets.keys() if v in misalignment_set_names
+    ]
+    default_dataset_name = default_dataset \
+        if default_dataset in datasets.keys() \
+        else ordinary_datasets[0] if len(ordinary_datasets) > 0 \
+        else misalignment_datasets[0] if len(misalignment_datasets) > 0 \
+        else ""
+    dataset_names = list(datasets.keys())
+
+    with gr.Row().style():
+        with gr.Column(scale=2):
+            with gr.Row():
+                dataset_dd = gr.Dropdown(dataset_names, label="Select dataset",
+                                         value="NODEFAULT", interactive=True)
+            with gr.Row():
+                disclaimer_ta = gr.Markdown(
+                    "## By clicking AGREE I consent to use the dataset "
+                    "for purely educational and academic purposes and "
+                    "not use it for any fraudulent activity; and I take "
+                    "all the responsibility if the data is used in a "
+                    "malicious application.", visible=False)
+            with gr.Row():
+                with gr.Column(scale=1):
+                    accept_disclaimer_bn = gr.Button("AGREE", visible=False)
+                with gr.Column(scale=1):
+                    decline_disclaimer_bn = gr.Button("DECLINE", visible=False)
+            with gr.Row():
+                with gr.Column(scale=3):
+                    assistant_dd = gr.Dropdown([], label="ASSISTANT", value="",
+                                               interactive=True)
+                with gr.Column(scale=3):
+                    user_dd = gr.Dropdown([], label="USER", value="",
+                                          interactive=True)
+        with gr.Column(scale=1):
+            gr.Markdown(
+                "## CAMEL: Communicative Agents for \"Mind\" Exploration"
+                " of Large Scale Language Model Society\n"
+                "Github repo: [https://github.com/lightaime/camel]"
+                "(https://github.com/lightaime/camel)\n"
+                '<div style="display:flex; justify-content:center;">'
+                '<img src="https://raw.githubusercontent.com/lightaime/camel/'
+                'master/misc/logo.png" alt="Logo" style="max-width:50%;">'
+                '</div>')
+
+    task_dd = gr.Dropdown([], label="Original task", value="",
+                          interactive=True)
+    specified_task_ta = gr.TextArea(label="Specified task", lines=2)
+    chatbot = gr.Chatbot()
+    accepted_st = gr.State(False)
+
+    def set_default_dataset() -> Dict:
+        """ Trigger for app load.
+
+        Returns:
+            Dict: Update dict for dataset_dd.
+        """
+        return gr.update(value=default_dataset_name)
+
+    def check_if_misalignment(dataset_name: str, accepted: bool) \
+            -> Tuple[Dict, Dict, Dict]:
+        """ Display AGREE/DECLINE if needed.
+
+        Returns:
+            Tuple: Visibility updates for the buttons.
+        """
+
+        if dataset_name == "misalignment" and not accepted:
+            return gr.update(visible=True), \
+                gr.update(visible=True), gr.update(visible=True)
+        else:
+            return gr.update(visible=False), \
+                gr.update(visible=False), gr.update(visible=False)
+
+    def enable_misalignment() -> Tuple[bool, Dict, Dict, Dict]:
+        """ Update the state of the accepted disclaimer.
+
+        Returns:
+            Tuple: New state and visibility updates for the buttons.
+        """
+
+        return True, gr.update(visible=False), \
+            gr.update(visible=False), gr.update(visible=False)
+
+    def disable_misalignment() -> Tuple[bool, Dict, Dict, Dict]:
+        """ Update the state of the accepted disclaimer.
+
+        Returns:
+            Tuple: New state and visibility updates for the buttons.
+        """
+
+        return False, gr.update(visible=False), \
+            gr.update(visible=False), gr.update(visible=False)
+
+    def update_dataset_selection(dataset_name: str,
+                                 accepted: bool) -> Tuple[Dict, Dict]:
+        """ Update roles based on the selected dataset.
+
+        Args:
+            dataset_name (str): Name of the loaded .zip dataset.
+            accepted (bool): If the disclaimer thas been accepted.
+
+        Returns:
+            Tuple[Dict, Dict]: New Assistant and User roles.
+        """
+
+        if dataset_name == "misalignment" and not accepted:
+            # If used did not accept the misalignment policy,
+            # keep the old selection.
+            return (gr.update(value="N/A",
+                              choices=[]), gr.update(value="N/A", choices=[]))
+
+        dataset = datasets[dataset_name]
+        assistant_roles = dataset['assistant_roles']
+        user_roles = dataset['user_roles']
+        assistant_role = random.choice(assistant_roles) \
+            if len(assistant_roles) > 0 else ""
+        user_role = random.choice(user_roles) if len(user_roles) > 0 else ""
+        return (gr.update(value=assistant_role, choices=assistant_roles),
+                gr.update(value=user_role, choices=user_roles))
+
+    def roles_dd_change(dataset_name: str, assistant_role: str,
+                        user_role: str) -> Dict:
+        """ Update the displayed chat upon inputs change.
+
+        Args:
+            assistant_role (str): Assistant dropdown value.
+            user_role (str): User dropdown value.
+
+        Returns:
+            Dict: New original roles state dictionary.
+        """
+        matrix = datasets[dataset_name]['matrix']
+        if (assistant_role, user_role) in matrix:
+            record: Dict[str, Dict] = matrix[(assistant_role, user_role)]
+            original_task_options = list(record.keys())
+            original_task = original_task_options[0]
+        else:
+            original_task = "N/A"
+            original_task_options = []
+
+        choices = gr.Dropdown.update(choices=original_task_options,
+                                     value=original_task, interactive=True)
+        return choices
+
+    def build_chat_history(messages: Dict[int, Dict]) -> List[Tuple]:
+        """ Structures chatbot contents from the loaded data.
+
+        Args:
+            messages (Dict[int, Dict]): Messages loaded from JSON.
+
+        Returns:
+            List[Tuple]: Chat history in chatbot UI element format.
+        """
+        history = []
+        curr_qa = (None, None)
+        for k in sorted(messages.keys()):
+            msg = messages[k]
+            content = msg['content']
+            if msg['role_type'] == "USER":
+                if curr_qa[0] is not None:
+                    history.append(curr_qa)
+                    curr_qa = (content, None)
+                else:
+                    curr_qa = (content, None)
+            elif msg['role_type'] == "ASSISTANT":
+                curr_qa = (curr_qa[0], content)
+                history.append(curr_qa)
+                curr_qa = (None, None)
+            else:
+                pass
+        return history
+
+    def task_dd_change(dataset_name: str, assistant_role: str, user_role: str,
+                       original_task: str) -> Tuple[str, List]:
+        """ Load task details and chatbot history into UI elements.
+
+        Args:
+            assistant_role (str): An assistan role.
+            user_role (str): An user role.
+            original_task (str): The original task.
+
+        Returns:
+            Tuple[str, List]: New contents of the specified task
+            and chatbot history UI elements.
+        """
+
+        matrix = datasets[dataset_name]['matrix']
+        if (assistant_role, user_role) in matrix:
+            task_dict: Dict[str, Dict] = matrix[(assistant_role, user_role)]
+            if original_task in task_dict:
+                chat = task_dict[original_task]
+                specified_task = chat['specified_task']
+                history = build_chat_history(chat['messages'])
+            else:
+                specified_task = "N/A"
+                history = []
+        else:
+            specified_task = "N/A"
+            history = []
+        return specified_task, history
+
+    dataset_dd.change(check_if_misalignment, [dataset_dd, accepted_st],
+                      [disclaimer_ta, accept_disclaimer_bn,
+                       decline_disclaimer_bn]) \
+              .then(update_dataset_selection,
+                    [dataset_dd, accepted_st],
+                    [assistant_dd, user_dd])
+
+    accept_disclaimer_bn.click(enable_misalignment, None, [
+        accepted_st, disclaimer_ta, accept_disclaimer_bn, decline_disclaimer_bn
+    ]) \
+        .then(update_dataset_selection,
+              [dataset_dd, accepted_st],
+              [assistant_dd, user_dd])
+
+    decline_disclaimer_bn.click(disable_misalignment, None, [
+        accepted_st, disclaimer_ta, accept_disclaimer_bn, decline_disclaimer_bn
+    ]) \
+        .then(update_dataset_selection,
+              [dataset_dd, accepted_st],
+              [assistant_dd, user_dd])
+
+    func_args = (roles_dd_change, [dataset_dd, assistant_dd, user_dd], task_dd)
+    assistant_dd.change(*func_args)
+    user_dd.change(*func_args)
+
+    task_dd.change(task_dd_change,
+                   [dataset_dd, assistant_dd, user_dd, task_dd],
+                   [specified_task_ta, chatbot])
+
+    blocks.load(set_default_dataset, None, dataset_dd)
+
+
+def construct_blocks(data_path: str, default_dataset: Optional[str]):
+    """ Construct Blocs app but do not launch it.
+
+    Args:
+        data_path (str): Path to the set of ZIP datasets.
+        default_dataset (Optional[str]): Name of the default dataset,
+            without extension.
+
+    Returns:
+        gr.Blocks: Blocks instance.
+    """
+
+    print("Loading the dataset...")
+    datasets = load_datasets(data_path)
+    print("Dataset is loaded")
+
+    print("Getting Data Explorer web server online...")
+
+    with gr.Blocks() as blocks:
+        construct_ui(blocks, datasets, default_dataset)
+
+    return blocks
+
+
+def main():
+    """ Entry point. """
+
+    args = parse_arguments()
+
+    blocks = construct_blocks(args.data_path, args.default_dataset)
+
+    blocks.queue(args.concurrency_count) \
+          .launch(share=args.share, inbrowser=args.inbrowser,
+                  server_name=args.server_name, server_port=args.server_port)
+
+    print("Exiting.")
+
+
+if __name__ == "__main__":
+    main()

apps/data_explorer/downloader.py

@@ -0,0 +1,42 @@
+import os
+import urllib.request
+
+from huggingface_hub import hf_hub_download
+
+REPO_ROOT = os.path.realpath(
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+
+
+def download_data():
+
+    print("Downloading...")
+
+    data_dir = os.path.join(REPO_ROOT, "datasets/")
+
+    os.makedirs(data_dir, exist_ok=True)
+
+    try:
+        hf_hub_download(repo_id="camel-ai/ai_society", repo_type="dataset",
+                        filename="ai_society_chat.zip", local_dir=data_dir,
+                        local_dir_use_symlinks=False)
+
+        hf_hub_download(repo_id="camel-ai/code", repo_type="dataset",
+                        filename="code_chat.zip", local_dir=data_dir,
+                        local_dir_use_symlinks=False)
+    except:
+        for name in ("ai_society_chat.zip", "code_chat.zip"):
+            data_url = ("https://storage.googleapis.com/"
+                        f"camel-bucket/datasets/private/{name}")
+            file_path = os.path.join(data_dir, os.path.split(data_url)[1])
+            urllib.request.urlretrieve(data_url, file_path)
+
+    data_url = ("https://storage.googleapis.com/"
+                "camel-bucket/datasets/private/misalignment.zip")
+    file_path = os.path.join(data_dir, os.path.split(data_url)[1])
+    urllib.request.urlretrieve(data_url, file_path)
+
+    print("Download done")
+
+
+if __name__ == "__main__":
+    download_data()

apps/data_explorer/loader.py

@@ -0,0 +1,172 @@
+"""
+Everything related to parsing the data JSONs into UI-compatible format.
+"""
+
+import glob
+import json
+import os
+import re
+import zipfile
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from tqdm import tqdm
+
+ChatHistory = Dict[str, Any]
+ParsedChatHistory = Dict[str, Any]
+AllChats = Dict[str, Any]
+Datasets = Dict[str, AllChats]
+
+REPO_ROOT = os.path.realpath(
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+
+
+class AutoZip:
+    def __init__(self, zip_path: str, ext: str = ".json"):
+        self.zip_path = zip_path
+        self.zip = zipfile.ZipFile(zip_path, "r")
+        self.fl = [f for f in self.zip.filelist if f.filename.endswith(ext)]
+
+    def __next__(self):
+        if self.index >= len(self.fl):
+            raise StopIteration
+        else:
+            finfo = self.fl[self.index]
+            with self.zip.open(finfo) as f:
+                raw_json = json.loads(f.read().decode("utf-8"))
+            self.index += 1
+            return raw_json
+
+    def __len__(self):
+        return len(self.fl)
+
+    def __iter__(self):
+        self.index = 0
+        return self
+
+
+def parse(raw_chat: ChatHistory) -> Union[ParsedChatHistory, None]:
+    """ Gets the JSON raw chat data, validates it and transforms
+        into an easy to work with form.
+
+    Args:
+        raw_chat (ChatHistory): In-memory loaded JSON data file.
+
+    Returns:
+        Union[ParsedChatHistory, None]: Parsed chat data or None
+        if there were parsing errors.
+    """
+
+    if "role_1" not in raw_chat:
+        return None
+
+    role_1 = raw_chat["role_1"]
+    if "_RoleType.ASSISTANT" not in role_1:
+        return None
+    assistant_role = role_1.split("_RoleType.ASSISTANT")
+    if len(assistant_role) < 1:
+        return None
+    if len(assistant_role[0]) <= 0:
+        return None
+    assistant_role = assistant_role[0]
+
+    role_2 = raw_chat["role_2"]
+    if "_RoleType.USER" not in role_2:
+        return None
+    user_role = role_2.split("_RoleType.USER")
+    if len(user_role) < 1:
+        return None
+    if len(user_role[0]) <= 0:
+        return None
+    user_role = user_role[0]
+
+    original_task = raw_chat["original_task"]
+    if len(original_task) <= 0:
+        return None
+
+    specified_task = raw_chat["specified_task"]
+    if len(specified_task) <= 0:
+        return None
+
+    messages = dict()
+    for key in raw_chat:
+        match = re.search("message_(?P<number>[0-9]+)", key)
+        if match:
+            number = int(match.group("number"))
+            messages[number] = raw_chat[key]
+
+    return dict(
+        assistant_role=assistant_role,
+        user_role=user_role,
+        original_task=original_task,
+        specified_task=specified_task,
+        messages=messages,
+    )
+
+
+def load_zip(zip_path: str) -> AllChats:
+    """ Load all JSONs from a zip file and parse them.
+
+    Args:
+        path (str): path to the ZIP file.
+
+    Returns:
+        AllChats: A dictionary with all possible assistant and
+        user roles and the matrix of chats.
+    """
+
+    zip_inst = AutoZip(zip_path)
+    parsed_list = []
+    for raw_chat in tqdm(iter(zip_inst)):
+        parsed = parse(raw_chat)
+        if parsed is None:
+            continue
+        parsed_list.append(parsed)
+
+    assistant_roles = set()
+    user_roles = set()
+    for parsed in parsed_list:
+        assistant_roles.add(parsed['assistant_role'])
+        user_roles.add(parsed['user_role'])
+    assistant_roles = list(sorted(assistant_roles))
+    user_roles = list(sorted(user_roles))
+    matrix: Dict[Tuple[str, str], List[Dict]] = dict()
+    for parsed in parsed_list:
+        key = (parsed['assistant_role'], parsed['user_role'])
+        original_task = parsed['original_task']
+        new_item = {
+            k: v
+            for k, v in parsed.items()
+            if k not in {'assistant_role', 'user_role', 'original_task'}
+        }
+        if key in matrix:
+            matrix[key][original_task] = new_item
+        else:
+            matrix[key] = {original_task: new_item}
+
+    return dict(
+        assistant_roles=assistant_roles,
+        user_roles=user_roles,
+        matrix=matrix,
+    )
+
+
+def load_datasets(path: Optional[str] = None) -> Datasets:
+    """ Load all JSONs from a set of zip files and parse them.
+
+    Args:
+        path (str): path to the folder with ZIP datasets.
+
+    Returns:
+        Datasets: A dictionary of dataset name and dataset contents.
+    """
+
+    if path is None:
+        path = os.path.join(REPO_ROOT, "datasets")
+
+    filt = os.path.join(path, "*.zip")
+    files = glob.glob(filt)
+    datasets = {}
+    for file_name in tqdm(files):
+        name = os.path.splitext(os.path.basename(file_name))[0]
+        datasets[name] = load_zip(file_name)
+    return datasets