Add a flagging callback to save json files to a hugging face dataset

gradio/__init__.py

@@ -48,6 +48,7 @@
     FlaggingCallback,
     HuggingFaceDatasetSaver,
     SimpleCSVLogger,
+    HuggingFaceDatasetJSONSaver,
 )
 from gradio.interface import Interface, TabbedInterface, close_all
 from gradio.ipython_ext import load_ipython_extension

gradio/flagging.py

@@ -5,6 +5,8 @@
 import io
 import json
 import os
+import random
+import string
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, List, Optional
 
@@ -344,3 +346,170 @@ def flag(
         self.repo.push_to_hub(commit_message="Flagged sample #{}".format(line_count))
 
         return line_count
+
+
+class HuggingFaceDatasetJSONSaver(FlaggingCallback):
+    """
+    A FlaggingCallback that saves flagged data to a HuggingFace dataset in JSONL format.
+    """
+
+    def __init__(
+        self,
+        hf_foken: str,
+        dataset_name: str,
+        organization: Optional[str] = None,
+        private: bool = False,
+        verbose: bool = True,
+    ):
+        """
+        Params:
+        hf_token (str): The token to use to access the huggingface API.
+        dataset_name (str): The name of the dataset to save the data to, e.g.
+            "image-classifier-1"
+        organization (str): The name of the organization to which to attach
+            the datasets. If None, the dataset attaches to the user only.
+        private (bool): If the dataset does not already exist, whether it
+            should be created as a private dataset or public. Private datasets
+            may require paid huggingface.co accounts
+        verbose (bool): Whether to print out the status of the dataset
+            creation.
+        """
+        self.hf_foken = hf_foken
+        self.dataset_name = dataset_name
+        self.organization_name = organization
+        self.dataset_private = private
+        self.verbose = verbose
+
+    def setup(self, components: List[Component], flagging_dir: str):
+        """
+        Params:
+        flagging_dir (str): local directory where the dataset is cloned,
+        updated, and pushed from.
+        """
+        try:
+            import huggingface_hub
+        except (ImportError, ModuleNotFoundError):
+            raise ImportError(
+                "Package `huggingface_hub` not found is needed "
+                "for HuggingFaceDatasetSaver. Try 'pip install huggingface_hub'."
+            )
+        path_to_dataset_repo = huggingface_hub.create_repo(
+            name=self.dataset_name,
+            token=self.hf_foken,
+            private=self.dataset_private,
+            repo_type="dataset",
+            exist_ok=True,
+        )
+        self.path_to_dataset_repo = path_to_dataset_repo  # e.g. "https://huggingface.co/datasets/abidlabs/test-audio-10"
+        self.components = components
+        self.flagging_dir = flagging_dir
+        self.dataset_dir = os.path.join(flagging_dir, self.dataset_name)
+        self.repo = huggingface_hub.Repository(
+            local_dir=self.dataset_dir,
+            clone_from=path_to_dataset_repo,
+            use_auth_token=self.hf_foken,
+        )
+        self.repo.git_pull()
+
+        self.infos_file = os.path.join(self.dataset_dir, "dataset_infos.json")
+
+    def flag(
+        self,
+        flag_data: List[Any],
+        flag_option: Optional[str] = None,
+        flag_index: Optional[int] = None,
+        username: Optional[str] = None,
+    ) -> int:
+        self.repo.git_pull(lfs=True)
+
+        # Generate unique folder for the flagged sample
+        unique_name = self.get_unique_name()  # unique name for folder
+        folder_name = os.path.join(
+            self.dataset_dir, unique_name
+        )  # unique folder for specific example
+        os.makedirs(folder_name, exist_ok=True)
+
+        # Now uses the existence of `dataset_infos.json` to determine if new
+        is_new = not os.path.exists(self.infos_file)
+
+        infos = {"flagged": {"features": {}}}
+
+        # File previews for certain input and output types
+        file_preview_types = {
+            gr.inputs.Audio: "Audio",
+            gr.outputs.Audio: "Audio",
+            gr.inputs.Image: "Image",
+            gr.outputs.Image: "Image",
+        }
+
+        # Generate the headers and dataset_infos
+        if is_new:
+
+            for component, sample in zip(self.components, flag_data):
+                infos["flagged"]["features"][component.label] = {
+                    "dtype": "string",
+                    "_type": "Value",
+                }
+                if isinstance(component, tuple(file_preview_types)):
+                    for _component, _type in file_preview_types.items():
+                        if isinstance(component, _component):
+                            infos["flagged"]["features"][component.label + " file"] = {
+                                "_type": _type
+                            }
+                            break
+
+            infos["flagged"]["features"]["flag"] = {
+                "dtype": "string",
+                "_type": "Value",
+            }
+
+        # Generate the row and header corresponding to the flagged sample
+        csv_data = []
+        headers = []
+
+        for component, sample in zip(self.components, flag_data):
+            headers.append(component.label)
+
+            try:
+                filepath = component.save_flagged(
+                    folder_name, component.label, sample, None
+                )
+            except Exception:
+                # Could not parse 'sample' (mostly) because it was None and `component.save_flagged` does not handle None cases. for example: Label (line 3109 of components.py raises an error if data is None)
+                filepath = None
+
+            if isinstance(component, tuple(file_preview_types)):
+                headers.append(component.label + " file")
+
+                csv_data.append(
+                    "{}/resolve/main/{}/{}".format(
+                        self.path_to_dataset_repo, unique_name, filepath
+                    )
+                    if filepath is not None
+                    else None
+                )
+
+            csv_data.append(filepath)
+        headers.append("flag")
+        csv_data.append(flag_option if flag_option is not None else "")
+
+        # Creates metadata dict from row data and dumps it
+        metadata_dict = {
+            header: _csv_data for header, _csv_data in zip(headers, csv_data)
+        }
+        self.dump_json(metadata_dict, os.path.join(folder_name, "metadata.jsonl"))
+
+        if is_new:
+            json.dump(infos, open(self.infos_file, "w"))
+
+        self.repo.push_to_hub(commit_message="Flagged sample {}".format(unique_name))
+        return unique_name
+
+    def get_unique_name(self):
+        return "".join(
+            [random.choice(string.ascii_letters + string.digits) for n in range(32)]
+        )
+
+    def dump_json(self, thing: dict, file_path: str) -> None:
+        with open(file_path, "w+", encoding="utf8") as f:
+            json.dump(thing, f)

test/test_flagging.py

@@ -67,6 +67,41 @@ def test_saver_flag(self):
             self.assertEqual(row_count, 2)  # 3 rows written including header
 
 
+class TestHuggingFaceDatasetJSONSaver(unittest.TestCase):
+    def test_saver_setup(self):
+        huggingface_hub.create_repo = MagicMock()
+        huggingface_hub.Repository = MagicMock()
+        flagger = flagging.HuggingFaceDatasetJSONSaver("test", "test")
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            flagger.setup([gr.Audio, gr.Textbox], tmpdirname)
+        huggingface_hub.create_repo.assert_called_once()
+
+    def test_saver_flag(self):
+        huggingface_hub.create_repo = MagicMock()
+        huggingface_hub.Repository = MagicMock()
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            io = gr.Interface(
+                lambda x: x,
+                "text",
+                "text",
+                flagging_dir=tmpdirname,
+                flagging_callback=flagging.HuggingFaceDatasetJSONSaver("test", "test"),
+            )
+            test_dir = os.path.join(tmpdirname, "test")
+            os.mkdir(test_dir)
+            io.launch(prevent_thread_lock=True)
+            row_unique_name = io.flagging_callback.flag(["test", "test"])
+            # Test existence of metadata.jsonl file for that example
+            self.assertEqual(
+                os.path.isfile(
+                    os.path.join(
+                        os.path.join(test_dir, row_unique_name), "metadata.jsonl"
+                    )
+                ),
+                True,
+            )
+
+
 class TestDisableFlagging(unittest.TestCase):
     def test_flagging_no_permission_error_with_flagging_disabled(self):
         with tempfile.TemporaryDirectory() as tmpdirname: