Update parquet scheduler

Co-authored-by: Lucain Pouget <[email protected]>

scheduler.py

@@ -16,6 +16,25 @@ class ParquetScheduler(CommitScheduler):
                 self.rows = []
             self.rows.append(row)
 
+    def set_schema(self, schema: Dict[str, Dict[str, str]]) -> None:
+        """
+        Define a schema to help `datasets` load the generated library.
+        This method is optional and can be called once just after the scheduler had been created. If it is not called,
+        the schema is automatically inferred before pushing the data to the Hub.
+        See https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Value for the list of
+        possible values.
+        Example:
+        ```py
+        scheduler.set_schema({
+            "prompt": {"_type": "Value", "dtype": "string"},
+            "negative_prompt": {"_type": "Value", "dtype": "string"},
+            "guidance_scale": {"_type": "Value", "dtype": "int64"},
+            "image": {"_type": "Image"},
+        })
+        ```
+        """
+        self._schema = schema
+
     def push_to_hub(self):
         # Check for new rows to push
         with self.lock:
@@ -25,29 +44,24 @@ class ParquetScheduler(CommitScheduler):
             return
 
         # Load images + create 'features' config for datasets library
-        hf_features: dict[str, Dict] = {}
+        hf_features: Dict[str, Dict] = getattr(self, '_schema', None) or {}
         path_to_cleanup: List[Path] = []
         for row in rows:
             for key, value in row.items():
-                if 'image' in key:
-                    # It's an image: we load the bytes, define a special schema and remember to cleanup the file
-                    # Note: could do the same with "Audio"
-                    image_path = Path(value)
-                    if image_path.is_file():
+                # Infer schema (for `datasets` library)
+                if key not in hf_features:
+                    hf_features[key] = _infer_schema(key, value)
+
+                # Load binary files if necessary
+                if hf_features[key]['_type'] in ('Image', 'Audio'):
+                    # It's an image or audio: we load the bytes and remember to cleanup the file
+                    file_path = Path(value)
+                    if file_path.is_file():
                         row[key] = {
-                            'path': image_path.name,
-                            'bytes': image_path.read_bytes()
-                        }
-                        path_to_cleanup.append(image_path)
-                        if key not in hf_features:
-                            hf_features[key] = {'_type': 'Image'}
-                else:
-                    # Otherwise, do nothing special
-                    if key not in hf_features:
-                        hf_features[key] = {
-                            '_type': 'Value',
-                            'dtype': 'string'
+                            'path': file_path.name,
+                            'bytes': file_path.read_bytes()
                         }
+                        path_to_cleanup.append(file_path)
 
         # Complete rows if needed
         for row in rows:
@@ -81,3 +95,24 @@ class ParquetScheduler(CommitScheduler):
         archive_file.close()
         for path in path_to_cleanup:
             path.unlink(missing_ok=True)
+
+
+def _infer_schema(key: str, value: Any) -> Dict[str, str]:
+    """Infer schema for the `datasets` library.
+
+    See https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Value.
+    """
+    if 'image' in key:
+        return {'_type': 'Image'}
+    if 'audio' in key:
+        return {'_type': 'Audio'}
+    if isinstance(value, int):
+        return {'_type': 'Value', 'dtype': 'int64'}
+    if isinstance(value, float):
+        return {'_type': 'Value', 'dtype': 'float64'}
+    if isinstance(value, bool):
+        return {'_type': 'Value', 'dtype': 'bool'}
+    if isinstance(value, bytes):
+        return {'_type': 'Value', 'dtype': 'binary'}
+    # Otherwise in last resort => convert it to a string
+    return {'_type': 'Value', 'dtype': 'string'}