first commit

app.py

@@ -0,0 +1,274 @@
+import gradio as gr
+
+from src import argilla_utils
+from src import dataset
+from src import spaces
+
+
+
+def refresh_dataset_settings_view(
+    columns,
+    question_columns,
+    field_columns,
+    split,
+    settings,
+    dataset_name,
+    argilla_dataset_name,
+    mapping,
+):
+    """This is a utility function to refresh the gradio applications state variables when a new dataset is loaded."""
+    columns = dataset.load_columns()
+    field_columns = dataset.get_field_columns()
+    question_columns = dataset.get_question_columns()
+    metadata_columns = []
+    vector_columns = []
+    split = dataset.load_split()
+    settings = None
+    dataset_name = dataset.load_dataset_name()
+    argilla_dataset_name = dataset.load_argilla_dataset_name()
+    mapping = None
+    return (
+        columns,
+        field_columns,
+        question_columns,
+        metadata_columns,
+        vector_columns,
+        split,
+        settings,
+        dataset_name,
+        argilla_dataset_name,
+        mapping,
+    )
+
+
+with gr.Blocks() as app:
+    ##############################################
+    # Define the app state
+    ##############################################
+
+    columns = gr.State(dataset.load_columns)
+    question_columns = gr.State(dataset.get_question_columns)
+    field_columns = gr.State(dataset.get_field_columns)
+    split = gr.State(dataset.load_split)
+    settings = gr.State(None)
+    dataset_name = gr.State(dataset.load_dataset_name)
+    argilla_dataset_name = gr.State(dataset.load_argilla_dataset_name)
+    mapping = gr.State(None)
+
+    state_variables = [
+        columns,
+        question_columns,
+        field_columns,
+        split,
+        settings,
+        dataset_name,
+        argilla_dataset_name,
+        mapping,
+    ]
+
+    ##############################################
+    # Define the app dataset and argilla space
+    ##############################################
+
+    gr.Markdown(
+        """# 🚂 Argilla Direct
+        A direct connection from a Hub dataset to an Argilla dataset.
+        This app allows you to create an Argilla dataset from a Hugging Face dataset. 
+        You will need to load a dataset from the Hugging Face Hub, create an Argilla space, 
+        define the dataset's settings, and add records to the dataset.
+        """
+    )
+
+    with gr.Group():
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column():
+                        dataset_name_input = gr.Textbox(
+                            label="Dataset Repo ID", value=dataset.load_dataset_name()
+                        )
+                    with gr.Column():
+                        split_input = gr.Dropdown(
+                            label="Dataset Split",
+                            choices=dataset.load_split_choices(),
+                            allow_custom_value=True,
+                            value=dataset.load_split(),
+                        )
+                    load_dataset_btn = gr.Button(value="1️⃣ Load Dataset")
+            with gr.Column():
+                argilla_space_name = gr.Textbox(
+                    label="Argilla Space Name", value=f"{dataset_name.value}_argilla"
+                )
+
+                create_argilla_space_btn = gr.Button(value="2️⃣ Create Argilla Space")
+
+    ##############################################
+    # Define the Argilla dataset configuration
+    ##############################################
+
+    gr.Markdown(
+        """## 3️⃣ Define Argilla Dataset
+        Define the settings for the Argilla dataset including fields, questions, metadata, and vectors.
+        Select the columns from the Hugging Face dataset to be used as Argilla dataset attributes.
+        """
+    )
+
+    with gr.Row():
+        with gr.Group():
+            with gr.Column():
+                # DATASET SETTINGS
+
+                # Argilla dataset name
+                argilla_dataset_name_view = gr.Textbox(
+                    label="Dataset Name",
+                    info="The name of the dataset in Argilla to be created or used",
+                    value=dataset.load_argilla_dataset_name(),
+                )
+                argilla_dataset_name_view.change(
+                    fn=lambda value: gr.update(
+                        value=dataset.load_argilla_dataset_name()
+                    ),
+                    inputs=[argilla_dataset_name_view],
+                    outputs=[argilla_dataset_name_view],
+                )
+
+                # Field columns
+                field_columns_view = gr.Dropdown(
+                    label="Field Columns",
+                    info="Columns to be used as fields in the Argilla dataset",
+                    choices=dataset.load_columns(),
+                    multiselect=True,
+                    value=dataset.get_field_columns(),
+                    allow_custom_value=True,
+                )
+                field_columns_view.change(
+                    fn=lambda value: gr.update(choices=dataset.load_columns()),
+                    inputs=[field_columns_view],
+                    outputs=[field_columns_view],
+                )
+
+                # Question columns
+                question_columns_view = gr.Dropdown(
+                    label="Question Columns",
+                    info="Columns to be used as question suggestions in the Argilla dataset",
+                    choices=dataset.load_columns(),
+                    multiselect=True,
+                    value=dataset.get_field_columns(),
+                    allow_custom_value=True,
+                )
+
+                question_columns_view.change(
+                    fn=lambda value: gr.update(choices=dataset.load_columns()),
+                    inputs=[question_columns_view],
+                    outputs=[question_columns_view],
+                )
+
+                with gr.Accordion(label="Define New Questions", open=False):
+                    with gr.Group():
+                        with gr.Column():
+                            question_type = gr.Dropdown(
+                                label="Question Type",
+                                info="The type of question to be added to the Argilla dataset",
+                                choices=["Text", "Label", "Rating"],
+                            )
+                        with gr.Column():
+                            question_name = gr.Textbox(
+                                label="Question Name",
+                                info="The name of the question to be added to the Argilla dataset",
+                            )
+                        with gr.Column():
+                            gr.Button(value="Add Question").click(
+                                fn=lambda type, name, questions: questions
+                                + [(type, name)],
+                                inputs=[
+                                    question_type,
+                                    question_name,
+                                    question_columns_view,
+                                ],
+                                outputs=[question_columns_view],
+                            )
+
+                with gr.Accordion(label="Define Metadata and Vectors", open=False):
+                    metadata_columns_view = gr.Dropdown(
+                        label="Metadata Columns",
+                        info="Columns to be used as metadata in the Argilla dataset",
+                        choices=dataset.load_columns(),
+                        multiselect=True,
+                    )
+                    vector_columns_view = gr.Dropdown(
+                        label="Vector Columns",
+                        info="Columns to be used as vectors in the Argilla dataset",
+                        choices=dataset.load_columns(),
+                        multiselect=True,
+                    )
+
+                n_records = gr.Slider(1, 10000, 100, label="Number of Records")
+                create_argilla_dataset_btn = gr.Button(value="Create Argilla Dataset")
+                add_records_btn = gr.Button(value="Add Records to Argilla")
+                delete_dataset_btn = gr.Button(value="Delete Argilla Dataset")
+
+        with gr.Column():
+            dataset_view = gr.Dataframe(
+                label="Dataset Viewer",
+                column_widths="20%",
+                headers=columns.value,
+                wrap=True,
+            )
+            records_view = gr.Text(label="Status", value="")
+
+    ##############################################
+    # Define the app logic
+    ##############################################
+
+    load_dataset_btn.click(
+        fn=dataset.load_dataset_from_hub,
+        inputs=[dataset_name_input],
+        outputs=[dataset_view],
+    ).then(
+        fn=refresh_dataset_settings_view,
+        inputs=state_variables,
+        outputs=[
+            columns,
+            question_columns_view,
+            field_columns_view,
+            split_input,
+            settings,
+            dataset_name,
+            argilla_dataset_name_view,
+            mapping,
+        ],
+    )
+
+    create_argilla_space_btn.click(
+        fn=spaces.create_argilla_space,
+        inputs=[argilla_space_name],
+        outputs=[records_view],
+    )
+
+    delete_dataset_btn.click(
+        fn=argilla_utils.delete_dataset,
+        inputs=[argilla_dataset_name_view],
+        outputs=[records_view],
+    )
+
+    create_argilla_dataset_btn.click(
+        fn=argilla_utils.define_dataset_setting,
+        inputs=[
+            argilla_dataset_name_view,
+            field_columns_view,
+            question_columns_view,
+            metadata_columns_view,
+            vector_columns_view,
+        ],
+        outputs=[records_view, mapping],
+    )
+
+    add_records_btn.click(
+        fn=argilla_utils.add_records,
+        inputs=[argilla_dataset_name_view, mapping, n_records],
+        outputs=[records_view],
+    )
+
+
+if __name__ == "__main__":
+    app.launch()

dataset_dir/dataset_dict.json

@@ -0,0 +1 @@
+{"splits": ["train"]}

dataset_dir/train/dataset_info.json

@@ -0,0 +1,94 @@
+{
+  "builder_name": "parquet",
+  "citation": "",
+  "config_name": "default",
+  "dataset_name": "10k_prompts_ranked",
+  "dataset_size": 8711680,
+  "description": "",
+  "download_checksums": {
+    "hf://datasets/DIBT/10k_prompts_ranked@3a9e44c398d92681e58b5c8ad39502203a002bac/data/train-00000-of-00001.parquet": {
+      "num_bytes": 3579688,
+      "checksum": null
+    }
+  },
+  "download_size": 3579688,
+  "features": {
+    "prompt": {
+      "dtype": "string",
+      "id": "field",
+      "_type": "Value"
+    },
+    "quality": [
+      {
+        "user_id": {
+          "dtype": "string",
+          "id": "question",
+          "_type": "Value"
+        },
+        "value": {
+          "dtype": "string",
+          "id": "suggestion",
+          "_type": "Value"
+        },
+        "status": {
+          "dtype": "string",
+          "id": "question",
+          "_type": "Value"
+        }
+      }
+    ],
+    "metadata": {
+      "dtype": "string",
+      "id": "metadata",
+      "_type": "Value"
+    },
+    "avg_rating": {
+      "dtype": "float64",
+      "_type": "Value"
+    },
+    "num_responses": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "agreement_ratio": {
+      "dtype": "float64",
+      "_type": "Value"
+    },
+    "raw_responses": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "kind": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "cluster_description": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "topic": {
+      "dtype": "string",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": "",
+  "size_in_bytes": 12291368,
+  "splits": {
+    "train": {
+      "name": "train",
+      "num_bytes": 8711680,
+      "num_examples": 10331,
+      "dataset_name": "10k_prompts_ranked"
+    }
+  },
+  "version": {
+    "version_str": "0.0.0",
+    "major": 0,
+    "minor": 0,
+    "patch": 0
+  }
+}

src/argilla_utils.py

@@ -0,0 +1,120 @@
+import argilla_sdk as rg
+from datasets import load_dataset
+
+from datasets import load_dataset
+
+from src.dataset import (
+    load_split,
+    is_label,
+    is_rating,
+    is_int,
+    is_float,
+    get_feature_values,
+    get_feature_labels,
+)
+
+client = rg.Argilla(api_url="http://localhost:6900", api_key="owner.apikey")
+
+
+def define_dataset_setting(
+    dataset_name, field_columns, question_columns, metadata_columns, vector_columns
+):
+    split = load_split()
+
+    fields, questions, metadata, vectors = [], [], [], []
+    mapping = {}
+
+    # Add field columns
+    for column_name in field_columns:
+        field_column_name = f"{column_name}_field"
+        fields.append(rg.TextField(name=field_column_name))
+        mapping[column_name] = field_column_name
+
+    # Add question columns
+    for column_name in question_columns:
+        if isinstance(column_name, (list, tuple)):
+            question_type, column_name = column_name
+        elif is_label(split, column_name):
+            question_type = "Label"
+        elif is_rating(split, column_name):
+            question_type = "Rating"
+        else:
+            question_type = "Text"
+
+        question_column_name = f"{column_name}_question"
+        if question_type == "Label":
+            values = get_feature_values(split, column_name)
+            titles = get_feature_labels(split, column_name)
+            labels = {str(l): feature for l, feature in zip(values, titles)}
+            questions.append(rg.LabelQuestion(name=question_column_name, labels=labels))
+        elif question_type == "Rating":
+            values = get_feature_values(split, column_name)
+            questions.append(
+                rg.RatingQuestion(name=question_column_name, values=values)
+            )
+        else:
+            questions.append(rg.TextQuestion(name=question_column_name))
+
+        if column_name in mapping:
+            column_name = f"{column_name}__"
+        mapping[column_name] = question_column_name
+
+    # Add metadata columns
+    if not metadata_columns:
+        metadata_columns = []
+
+    for column_name in metadata_columns:
+        metadata_column_name = f"{column_name}_metadata"
+        if is_int(split, column_name):
+            metadata.append(rg.IntegerMetadataProperty(name=metadata_column_name))
+        elif is_float(split, column_name):
+            metadata.append(rg.FloatMetadataProperty(name=metadata_column_name))
+        elif is_label:
+            values = list(map(str, get_feature_values(split, column_name)))
+            metadata.append(
+                rg.TermsMetadataProperty(name=metadata_column_name, options=values)
+            )
+        mapping[column_name] = metadata_column_name
+
+    # Add vector columns
+    if not vector_columns:
+        vector_columns = []
+
+    for column_name in vector_columns:
+        vectors.append(rg.VectorField(name=column_name))
+
+    settings = rg.Settings(
+        fields=fields, questions=questions, metadata=metadata, vectors=vectors
+    )
+
+    dataset = rg.Dataset(name=dataset_name, settings=settings, client=client)
+
+    if not dataset.exists():
+        dataset.create()
+
+    return str(settings.serialize()), mapping
+
+
+def add_records(argilla_dataset_name, mapping, n_records):
+    split = load_split()
+    df = load_dataset(load_repo_id())[split].take(n_records).to_pandas()
+    dataset = client.datasets(argilla_dataset_name)
+    questions = dataset.settings.questions
+    for question in questions:
+        if question.name in mapping.values():
+            column_name = [k for k, v in mapping.items() if v == question.name][0]
+            column_name = column_name.replace("__", "")
+            if is_label(split, column_name):
+                df[column_name] = df[column_name].apply(str)
+    for source, target in mapping.items():
+        if source.endswith("__"):
+            df[source] = df[source.replace("__", "")]
+    records = df.to_dict(orient="records")
+    dataset.records.log(records, mapping=mapping)
+    return f"{len(df)} records added with mapping {mapping}"
+
+
+def delete_dataset(argilla_dataset_name):
+    dataset = client.datasets(argilla_dataset_name)
+    dataset.delete()
+    return f"Dataset {argilla_dataset_name} deleted"

src/constants.py

@@ -0,0 +1,2 @@
+LOCAL_DATASET_PATH = "dataset_dir"
+LOCAL_CONFIG_PATH = "config_dir"

src/dataset.py

@@ -0,0 +1,159 @@
+import os
+import json
+
+from datetime import datetime
+
+from datasets import load_dataset
+
+from src.constants import LOCAL_CONFIG_PATH, LOCAL_DATASET_PATH
+
+
+##############################################
+# Get the dataset app
+##############################################
+
+
+def load_dataset_from_hub(dataset_name):
+    # delete the existing dataset
+    if os.path.exists(LOCAL_DATASET_PATH):
+        os.system(f"rm -rf {LOCAL_DATASET_PATH}")
+    ds = load_dataset(dataset_name)
+    ds.save_to_disk(LOCAL_DATASET_PATH)
+    split = load_split()
+    columns = list(ds[split].features.keys())
+    df = ds[split].to_pandas()
+    with open(LOCAL_CONFIG_PATH, "w") as f:
+        json.dump({"columns": columns, "split": split, "name": dataset_name}, f)
+    return df
+
+
+##############################################
+# Define the dataset app
+##############################################
+
+
+def load_repo_id():
+    with open(LOCAL_CONFIG_PATH, "r") as f:
+        config = json.load(f)
+    return config["name"]
+
+
+def load_dataset_dict_json(split):
+    dataset_dict_fn = "dataset_info.json"
+    path = os.path.join(LOCAL_DATASET_PATH, split, dataset_dict_fn)
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def load_dataset_name():
+    dataset_dict = load_dataset_dict_json("train")
+    return dataset_dict["dataset_name"]
+
+
+def load_argilla_dataset_name():
+    name = load_dataset_name()
+    now = datetime.now()
+    name = f"{name}_{now.strftime('%Y%m%d%H%M%S')}"
+    return name
+
+
+def load_split_choices():
+    dataset_dict = load_dataset_dict_json("train")
+    return list(dataset_dict["splits"].keys())
+
+
+def load_split():
+    return load_split_choices()[0]
+
+
+def load_columns():
+    dataset_dict = load_dataset_dict_json("train")
+    return list(dataset_dict["features"].keys())
+
+
+def get_split_features(split):
+    dataset_dict = load_dataset_dict_json(split)
+    return dataset_dict["features"]
+
+
+def get_feature_type(split, column_name):
+    features = get_split_features(split)
+    return features[column_name]["_type"]
+
+
+def get_feature_dtype(split, column_name):
+    features = get_split_features(split)
+    try:
+        return features[column_name]["dtype"]
+    except TypeError:
+        return None
+
+
+def is_field(split, column_name):
+    try:
+        return (
+            get_feature_dtype(split, column_name) == "string"
+            and get_feature_type(split, column_name) == "Value"
+        )
+    except KeyError:
+        return False
+
+
+def is_label(split, column_name):
+    feature_type = get_feature_type(split, column_name)
+    return feature_type == "ClassLabel"
+
+
+def is_float(split, column_name):
+    try:
+        feature_type = get_feature_type(split, column_name)
+        feature_dtype = get_feature_dtype(split, column_name)
+        return feature_type == "Value" and feature_dtype.startswith("float")
+    except KeyError:
+        return False
+
+
+def is_int(split, column_name):
+    try:
+        feature_type = get_feature_type(split, column_name)
+        feature_dtype = get_feature_dtype(split, column_name)
+        return feature_type == "Value" and feature_dtype.startswith("int")
+    except KeyError:
+        return False
+
+
+def get_feature_labels(split, column_name):
+    features = get_split_features(split)
+    return features[column_name]["names"]
+
+
+def get_feature_values(split, column_name):
+    ds = load_dataset(load_repo_id())
+    return list(set(ds[split][column_name]))
+
+
+def is_rating(split, column_name):
+    feature_values = get_feature_values(split, column_name)
+    if not is_int(split, column_name):
+        return False
+    if len(feature_values) > 10:
+        return False
+    return True
+
+
+def get_field_columns():
+    split = load_split()
+    columns = load_columns()
+    return [column for column in columns if is_field(split, column)]
+
+
+def get_question_columns():
+    split = load_split()
+    columns = load_columns()
+    return [column for column in columns if not is_field(split, column)]
+
+
+def load_dataset_df():
+    split = load_split()
+    ds = load_dataset(load_repo_id())
+    return ds[split].to_pandas()

src/spaces.py

@@ -0,0 +1,11 @@
+from huggingface_hub import duplicate_space
+
+
+def create_argilla_space(target_argilla_space):
+    duplicate_space(
+        from_id="argilla/argilla-template-space",
+        to_id=target_argilla_space,
+        private=False,
+        exist_ok=True,
+    )
+    return target_argilla_space