[FEAT] MIDI feature support

setup.py

@@ -145,6 +145,10 @@
     "Pillow>=9.4.0",  # When PIL.Image.ExifTags was introduced
 ]
 
+MIDI_REQUIRE = [
+    "pretty-midi>=0.2.0",
+]
+
 BENCHMARKS_REQUIRE = [
     "tensorflow==2.12.0",
     "torch==2.0.1",
@@ -187,6 +191,7 @@
     "Pillow>=9.4.0",  # When PIL.Image.ExifTags was introduced
     "torchcodec>=0.7.0; python_version < '3.14'",  # minium version to get windows support, torchcodec doesn't have wheels for 3.14 yet
     "nibabel>=5.3.1",
+    "pretty-midi>=0.2.0",
 ]
 
 NUMPY2_INCOMPATIBLE_LIBRARIES = [
@@ -213,6 +218,7 @@
 EXTRAS_REQUIRE = {
     "audio": AUDIO_REQUIRE,
     "vision": VISION_REQUIRE,
+    "midi": MIDI_REQUIRE,
     "tensorflow": [
         "tensorflow>=2.6.0",
     ],

src/datasets/arrow_dataset.py

@@ -78,7 +78,7 @@
 from .arrow_writer import ArrowWriter, OptimizedTypedSequence
 from .data_files import sanitize_patterns
 from .download.streaming_download_manager import xgetsize
-from .features import Audio, ClassLabel, Features, Image, List, Value, Video
+from .features import Audio, ClassLabel, Features, Image, List, Midi, Value, Video
 from .features.features import (
     FeatureType,
     _align_features,
@@ -5358,7 +5358,7 @@ def _estimate_nbytes(self) -> int:
 
             def extra_nbytes_visitor(array, feature):
                 nonlocal extra_nbytes
-                if isinstance(feature, (Audio, Image, Video)):
+                if isinstance(feature, (Audio, Image, Video, Midi)):
                     for x in array.to_pylist():
                         if x is not None and x["bytes"] is None and x["path"] is not None:
                             size = xgetsize(x["path"])

src/datasets/arrow_writer.py

@@ -25,7 +25,7 @@
 from fsspec.core import url_to_fs
 
 from . import config
-from .features import Audio, Features, Image, Pdf, Value, Video
+from .features import Audio, Features, Image, Midi, Pdf, Value, Video
 from .features.features import (
     FeatureType,
     List,
@@ -78,6 +78,8 @@ def set_batch_size(feature: FeatureType) -> None:
             batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_AUDIO_DATASETS)
         elif isinstance(feature, Video) and config.ARROW_RECORD_BATCH_SIZE_FOR_VIDEO_DATASETS is not None:
             batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_VIDEO_DATASETS)
+        elif isinstance(feature, Midi) and config.ARROW_RECORD_BATCH_SIZE_FOR_MIDI_DATASETS is not None:
+            batch_size = min(batch_size, config.ARROW_RECORD_BATCH_SIZE_FOR_MIDI_DATASETS)
         elif (
             isinstance(feature, Value)
             and feature.dtype == "binary"
@@ -118,6 +120,8 @@ def set_batch_size(feature: FeatureType) -> None:
             batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS)
         elif isinstance(feature, Video) and config.PARQUET_ROW_GROUP_SIZE_FOR_VIDEO_DATASETS is not None:
             batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_VIDEO_DATASETS)
+        elif isinstance(feature, Midi) and config.PARQUET_ROW_GROUP_SIZE_FOR_MIDI_DATASETS is not None:
+            batch_size = min(batch_size, config.PARQUET_ROW_GROUP_SIZE_FOR_MIDI_DATASETS)
         elif (
             isinstance(feature, Value)
             and feature.dtype == "binary"

src/datasets/config.py

@@ -198,12 +198,14 @@
 PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS = None
 PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS = None
 PARQUET_ROW_GROUP_SIZE_FOR_VIDEO_DATASETS = None
+PARQUET_ROW_GROUP_SIZE_FOR_MIDI_DATASETS = None
 
 # Arrow configuration
 ARROW_RECORD_BATCH_SIZE_FOR_AUDIO_DATASETS = 100
 ARROW_RECORD_BATCH_SIZE_FOR_IMAGE_DATASETS = 100
 ARROW_RECORD_BATCH_SIZE_FOR_BINARY_DATASETS = 100
 ARROW_RECORD_BATCH_SIZE_FOR_VIDEO_DATASETS = 10
+ARROW_RECORD_BATCH_SIZE_FOR_MIDI_DATASETS = 100
 
 # Offline mode
 _offline = os.environ.get("HF_DATASETS_OFFLINE")

src/datasets/features/__init__.py

@@ -16,10 +16,12 @@
     "Video",
     "Pdf",
     "Nifti",
+    "Midi",
 ]
 from .audio import Audio
 from .features import Array2D, Array3D, Array4D, Array5D, ClassLabel, Features, LargeList, List, Sequence, Value
 from .image import Image
+from .midi import Midi
 from .nifti import Nifti
 from .pdf import Pdf
 from .translation import Translation, TranslationVariableLanguages

src/datasets/features/features.py

@@ -42,6 +42,7 @@
 from ..utils.py_utils import asdict, first_non_null_value, zip_dict
 from .audio import Audio
 from .image import Image, encode_pil_image
+from .midi import Midi
 from .nifti import Nifti
 from .pdf import Pdf, encode_pdfplumber_pdf
 from .translation import Translation, TranslationVariableLanguages
@@ -1431,6 +1432,7 @@ def decode_nested_example(schema, obj, token_per_repo_id: Optional[dict[str, Uni
     Video.__name__: Video,
     Pdf.__name__: Pdf,
     Nifti.__name__: Nifti,
+    Midi.__name__: Midi,
 }