Repository: huggingface/dataset-viewer

Issue:
Avoid unnecessary conversions of audio files without extension
This will improve the speed performance of the /rows endpoint.

As we analogously do in the `datasets` library and suggested by @lhoestq, we could infer the file format from the header of the file (magic bytes).

To be done, once this PR is merged:
- #2180


Discussion:

Resolving patch:
diff --git a/libs/libcommon/src/libcommon/viewer_utils/asset.py b/libs/libcommon/src/libcommon/viewer_utils/asset.py
--- a/libs/libcommon/src/libcommon/viewer_utils/asset.py
+++ b/libs/libcommon/src/libcommon/viewer_utils/asset.py
@@ -4,7 +4,7 @@
 from io import BytesIO
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import TypedDict
+from typing import Optional, TypedDict
 from urllib import parse
 
 from PIL import Image  # type: ignore
@@ -86,7 +86,7 @@ def create_audio_file(
     row_idx: int,
     column: str,
     audio_file_bytes: bytes,
-    audio_file_extension: str,
+    audio_file_extension: Optional[str],
     filename: str,
     public_assets_storage: PublicAssetsStorage,
 ) -> list[AudioSource]:
@@ -121,7 +121,7 @@ def create_audio_file(
                 f.write(audio_file_bytes)
         else:  # we need to convert
             # might spawn a process to convert the audio file using ffmpeg
-            with NamedTemporaryFile("wb", suffix=audio_file_extension or None) as tmpfile:
+            with NamedTemporaryFile("wb", suffix=audio_file_extension) as tmpfile:
                 tmpfile.write(audio_file_bytes)
                 segment: AudioSegment = AudioSegment.from_file(
                     tmpfile.name, audio_file_extension[1:] if audio_file_extension else None
diff --git a/libs/libcommon/src/libcommon/viewer_utils/features.py b/libs/libcommon/src/libcommon/viewer_utils/features.py
--- a/libs/libcommon/src/libcommon/viewer_utils/features.py
+++ b/libs/libcommon/src/libcommon/viewer_utils/features.py
@@ -31,6 +31,10 @@
 from libcommon.viewer_utils.asset import create_audio_file, create_image_file
 
 UNSUPPORTED_FEATURES = [Value("binary")]
+AUDIO_FILE_MAGIC_NUMBERS: dict[str, Any] = {
+    ".wav": [(b"\x52\x49\x46\x46", 0), (b"\x57\x41\x56\x45", 8)],  # AND: (magic_number, start)
+    ".mp3": (b"\xFF\xFB", b"\xFF\xF3", b"\xFF\xF2", b"\x49\x44\x33"),  # OR
+}
 
 
 def append_hash_suffix(string: str, json_path: Optional[list[Union[str, int]]] = None) -> str:
@@ -118,19 +122,28 @@ def audio(
             "Audio cell must be an encoded dict of an audio sample, "
             f"but got {str(value)[:300]}{'...' if len(str(value)) > 300 else ''}"
         )
+    audio_file_extension = get_audio_file_extension(value)
+    audio_file_bytes = get_audio_file_bytes(value)
+    if not audio_file_extension:
+        audio_file_extension = infer_audio_file_extension(audio_file_bytes)
+    # convert to wav if the file is not wav or mp3 already
+    target_audio_file_extension = audio_file_extension if audio_file_extension in [".wav", ".mp3"] else ".wav"
+    # this function can raise, we don't catch it
+    return create_audio_file(
+        dataset=dataset,
+        revision=revision,
+        config=config,
+        split=split,
+        row_idx=row_idx,
+        column=featureName,
+        audio_file_bytes=audio_file_bytes,
+        audio_file_extension=audio_file_extension,
+        public_assets_storage=public_assets_storage,
+        filename=f"{append_hash_suffix('audio', json_path)}{target_audio_file_extension}",
+    )
 
-    if "path" in value and isinstance(value["path"], str):
-        # .split("::")[0] for chained URLs like zip://audio.wav::https://foo.bar/data.zip
-        # It might be "" for audio files downloaded from the Hub
-        audio_file_extension = os.path.splitext(value["path"].split("::")[0])[1]
-    elif ("path" in value and value["path"] is None) or "array" in value:
-        audio_file_extension = ".wav"
-    else:
-        raise ValueError(
-            "An audio sample should have 'path' and 'bytes' (or 'array' and 'sampling_rate') but got"
-            f" {', '.join(value)}."
-        )
 
+def get_audio_file_bytes(value: Any) -> bytes:
     if "bytes" in value and isinstance(value["bytes"], bytes):
         audio_file_bytes = value["bytes"]
     elif "path" in value and isinstance(value["path"], str) and os.path.exists(value["path"]):
@@ -150,23 +163,33 @@ def audio(
             "An audio sample should have 'path' and 'bytes' (or 'array' and 'sampling_rate') but got"
             f" {', '.join(value)}."
         )
+    return audio_file_bytes
 
-    # convert to wav if the file is not wav or mp3 already
-    ext = audio_file_extension if audio_file_extension in [".wav", ".mp3"] else ".wav"
 
-    # this function can raise, we don't catch it
-    return create_audio_file(
-        dataset=dataset,
-        revision=revision,
-        config=config,
-        split=split,
-        row_idx=row_idx,
-        column=featureName,
-        audio_file_bytes=audio_file_bytes,
-        audio_file_extension=audio_file_extension,
-        public_assets_storage=public_assets_storage,
-        filename=f"{append_hash_suffix('audio', json_path)}{ext}",
-    )
+def get_audio_file_extension(value: Any) -> Optional[str]:
+    if "path" in value and isinstance(value["path"], str):
+        # .split("::")[0] for chained URLs like zip://audio.wav::https://foo.bar/data.zip
+        # It might be "" for audio files downloaded from the Hub: make it None
+        audio_file_extension = os.path.splitext(value["path"].split("::")[0])[1] or None
+    elif ("path" in value and value["path"] is None) or "array" in value:
+        audio_file_extension = ".wav"
+    else:
+        raise ValueError(
+            "An audio sample should have 'path' and 'bytes' (or 'array' and 'sampling_rate') but got"
+            f" {', '.join(value)}."
+        )
+    return audio_file_extension
+
+
+def infer_audio_file_extension(audio_file_bytes: bytes) -> Optional[str]:
+    for audio_file_extension, magic_numbers in AUDIO_FILE_MAGIC_NUMBERS.items():
+        if isinstance(magic_numbers, list):
+            if all(audio_file_bytes.startswith(magic_number, start) for magic_number, start in magic_numbers):
+                return audio_file_extension
+        else:
+            if audio_file_bytes.startswith(magic_numbers):
+                return audio_file_extension
+    return None
 
 
 def get_cell_value(


Test patch:
diff --git a/libs/libcommon/tests/viewer_utils/data/test_audio_16000.mp3 b/libs/libcommon/tests/viewer_utils/data/test_audio_16000.mp3
new file mode 100644
Binary files /dev/null and b/libs/libcommon/tests/viewer_utils/data/test_audio_16000.mp3 differ
diff --git a/libs/libcommon/tests/viewer_utils/test_features.py b/libs/libcommon/tests/viewer_utils/test_features.py
--- a/libs/libcommon/tests/viewer_utils/test_features.py
+++ b/libs/libcommon/tests/viewer_utils/test_features.py
@@ -22,6 +22,7 @@
 from libcommon.viewer_utils.features import (
     get_cell_value,
     get_supported_unsupported_columns,
+    infer_audio_file_extension,
 )
 
 ASSETS_FOLDER = "assets"
@@ -456,3 +457,18 @@ async def convert_to_response_dict(http_response, operation_model):  # type: ign
                     "type": "audio/wav",
                 },
             ]
+
+
+@pytest.mark.parametrize(
+    "audio_file_name, expected_audio_file_extension",
+    [("test_audio_44100.wav", ".wav"), ("test_audio_16000.mp3", ".mp3")],
+)
+def test_infer_audio_file_extension(
+    audio_file_name: str,
+    expected_audio_file_extension: str,
+    shared_datadir: Path,
+    public_assets_storage: PublicAssetsStorage,
+) -> None:
+    audio_file_bytes = (shared_datadir / audio_file_name).read_bytes()
+    audio_file_extension = infer_audio_file_extension(audio_file_bytes)
+    assert audio_file_extension == expected_audio_file_extension