Apply fix to the HF dataset saver

app.py

@@ -7,15 +7,16 @@ from difflib import ndiff
 import gradio as gr
 
 from data_loader import load_data
+from hf_dataset_saver_builder import get_dataset_saver
 
-HF_TOKEN = os.environ.get('HF_TOKEN')
-HF_DATASET = os.environ.get('HF_DATASET')
+HF_TOKEN = os.environ.get('HF_REWRITING_TOKEN')
+HF_DATASET = os.environ.get('HF_REWRITING_DATASET')
 
 data = load_data()
 
 n_samples = len(data)
 
-saver = gr.HuggingFaceDatasetSaver(HF_TOKEN, HF_DATASET, private=True, separate_dirs=True)
+saver = get_dataset_saver(HF_TOKEN, HF_DATASET, private=True, separate_dirs=True)
 
 
 def convert_diff_to_unified(diff):

hf_dataset_saver_builder.py

@@ -0,0 +1,71 @@
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+
+
+def _deserialize_components_fix(
+        self,
+        data_dir: Path,
+        flag_data: list[Any],
+        flag_option: str = "",
+        username: str = "",
+) -> tuple[dict[Any, Any], list[Any]]:
+    """Deserialize components and return the corresponding row for the flagged sample.
+
+    Images/audio are saved to disk as individual files.
+    """
+    # Components that can have a preview on dataset repos
+    file_preview_types = {gr.Audio: "Audio", gr.Image: "Image"}
+
+    # Generate the row corresponding to the flagged sample
+    features = OrderedDict()
+    row = []
+    for component, sample in zip(self.components, flag_data):
+        # Get deserialized object (will save sample to disk if applicable -file, audio, image,...-)
+        label = component.label or ""
+        save_dir = data_dir / gr.flagging.client_utils.strip_invalid_filename_characters(label)
+        save_dir.mkdir(exist_ok=True, parents=True)
+        deserialized = component.flag(sample, save_dir)
+
+        # Add deserialized object to row
+        features[label] = {"dtype": "string", "_type": "Value"}
+        try:
+            assert Path(deserialized).exists()
+            row.append(str(Path(deserialized).relative_to(self.dataset_dir)))
+        except (AssertionError, TypeError, ValueError, OSError):
+            deserialized = "" if deserialized is None else str(deserialized)
+            row.append(deserialized)
+
+        # If component is eligible for a preview, add the URL of the file
+        # Be mindful that images and audio can be None
+        if isinstance(component, tuple(file_preview_types)):  # type: ignore
+            for _component, _type in file_preview_types.items():
+                if isinstance(component, _component):
+                    features[label + " file"] = {"_type": _type}
+                    break
+            if deserialized:
+                path_in_repo = str(  # returned filepath is absolute, we want it relative to compute URL
+                    Path(deserialized).relative_to(self.dataset_dir)
+                ).replace("\\", "/")
+                row.append(
+                    gr.flagging.huggingface_hub.hf_hub_url(
+                        repo_id=self.dataset_id,
+                        filename=path_in_repo,
+                        repo_type="dataset",
+                    )
+                )
+            else:
+                row.append("")
+    features["flag"] = {"dtype": "string", "_type": "Value"}
+    features["username"] = {"dtype": "string", "_type": "Value"}
+    row.append(flag_option)
+    row.append(username)
+    return features, row
+
+
+def get_dataset_saver(*args, **kwargs):
+    saver = gr.HuggingFaceDatasetSaver(*args, **kwargs)
+    saver._deserialize_components = _deserialize_components_fix
+    return saver