Don't copy image/audio/video/pdf from repo file, and link to file directly in src (#3330)

* read from hf:// paths when needed

* fix source url

Author: lhoestq

libs/libapi/src/libapi/response.py

@@ -17,6 +17,8 @@ async def create_response(
     config: str,
     split: str,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     pa_table: pa.Table,
     offset: int,
     features: Features,
@@ -39,6 +41,8 @@ async def create_response(
             config=config,
             split=split,
             storage_client=storage_client,
+            hf_endpoint=hf_endpoint,
+            hf_token=hf_token,
             offset=offset,
             features=features,
             row_idx_column=ROW_IDX_COLUMN if use_row_idx_column else None,

libs/libapi/src/libapi/rows_utils.py

@@ -25,6 +25,8 @@ def _transform_row(
     split: str,
     features: Features,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     offset: int,
     row_idx_column: Optional[str],
 ) -> Row:
@@ -42,6 +44,8 @@ def _transform_row(
                 featureName=featureName,
                 fieldType=fieldType,
                 storage_client=storage_client,
+                hf_endpoint=hf_endpoint,
+                hf_token=hf_token,
             )
         except Exception as err:
             suggestion_messages: dict[type[Exception], str] = {
@@ -65,6 +69,8 @@ async def transform_rows(
     rows: list[Row],
     features: Features,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     offset: int,
     row_idx_column: Optional[str],
 ) -> list[Row]:
@@ -76,6 +82,8 @@ async def transform_rows(
         split=split,
         features=features,
         storage_client=storage_client,
+        hf_endpoint=hf_endpoint,
+        hf_token=hf_token,
         offset=offset,
         row_idx_column=row_idx_column,
     )

libs/libapi/src/libapi/utils.py

@@ -205,6 +205,8 @@ async def to_rows_list(
     offset: int,
     features: Features,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     row_idx_column: Optional[str] = None,
     truncated_columns: Optional[list[str]] = None,
 ) -> list[RowItem]:
@@ -217,6 +219,8 @@ async def to_rows_list(
         rows=pa_table.to_pylist(),
         features=features,
         storage_client=storage_client,
+        hf_endpoint=hf_endpoint,
+        hf_token=hf_token,
         offset=offset,
         row_idx_column=row_idx_column,
     )

libs/libapi/tests/test_response.py

@@ -39,6 +39,8 @@ async def test_create_response(storage_client: StorageClient) -> None:
         config="default",
         split="train",
         storage_client=storage_client,
+        hf_endpoint="not-needed",
+        hf_token=None,
         pa_table=ds.data,
         offset=0,
         features=ds.features,
@@ -63,6 +65,8 @@ async def test_create_response_with_row_idx_column(storage_client: StorageClient
         config="default",
         split="train",
         storage_client=storage_client,
+        hf_endpoint="not-needed",
+        hf_token=None,
         pa_table=ds.data,
         offset=0,
         features=ds.features,
@@ -91,6 +95,8 @@ async def test_create_response_with_image(image_path: str, storage_client: Stora
         config=config,
         split=split,
         storage_client=storage_client,
+        hf_endpoint="not-needed",
+        hf_token=None,
         pa_table=ds_image.data,
         offset=0,
         features=ds_image.features,
@@ -131,6 +137,8 @@ async def test_create_response_with_document(document_path: str, storage_client:
         config=config,
         split=split,
         storage_client=storage_client,
+        hf_endpoint="not-needed",
+        hf_token=None,
         pa_table=ds_document.data,
         offset=0,
         features=ds_document.features,

libs/libcommon/src/libcommon/viewer_utils/asset.py

@@ -8,14 +8,19 @@
 
 import fitz
 from pdfplumber.pdf import PDF
-from PIL import Image, ImageOps
+from PIL import Image
 from pydub import AudioSegment  # type:ignore
 
 if TYPE_CHECKING:
     from libcommon.storage_client import StorageClient
 
 
-SUPPORTED_AUDIO_EXTENSION_TO_MEDIA_TYPE = {".wav": "audio/wav", ".mp3": "audio/mpeg", ".opus": "audio/ogg"}
+SUPPORTED_AUDIO_EXTENSION_TO_MEDIA_TYPE = {
+    ".wav": "audio/wav",
+    ".mp3": "audio/mpeg",
+    ".opus": "audio/ogg",
+    ".flac": "audio/x-flac",
+}
 SUPPORTED_AUDIO_EXTENSIONS = SUPPORTED_AUDIO_EXTENSION_TO_MEDIA_TYPE.keys()
 DATASET_GIT_REVISION_PLACEHOLDER = "{dataset_git_revision}"
 LOCK = threading.Lock()
@@ -71,7 +76,6 @@ def create_image_file(
     )
     path = replace_dataset_git_revision_placeholder(object_path, revision=revision)
     if storage_client.overwrite or not storage_client.exists(path):
-        image = ImageOps.exif_transpose(image)  # type: ignore[assignment]
         buffer = BytesIO()
         image.save(fp=buffer, format=format)
         buffer.seek(0)

libs/libcommon/src/libcommon/viewer_utils/features.py

@@ -30,12 +30,17 @@
     Value,
     Video,
 )
+from huggingface_hub import HfFileSystem
 from PIL import Image as PILImage
 
 from libcommon.dtos import FeatureItem
 from libcommon.storage_client import StorageClient
 from libcommon.viewer_utils.asset import (
+    SUPPORTED_AUDIO_EXTENSION_TO_MEDIA_TYPE,
     SUPPORTED_AUDIO_EXTENSIONS,
+    AudioSource,
+    ImageSource,
+    VideoSource,
     create_audio_file,
     create_image_file,
     create_pdf_file,
@@ -89,6 +94,8 @@ def image(
     value: Any,
     featureName: str,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     json_path: Optional[list[Union[str, int]]] = None,
 ) -> Any:
     if value is None:
@@ -97,13 +104,15 @@ def image(
         value = PILImage.open(BytesIO(value["bytes"]))
     elif isinstance(value, bytes):
         value = PILImage.open(BytesIO(value))
-    elif (
-        isinstance(value, dict)
-        and "path" in value
-        and isinstance(value["path"], str)
-        and os.path.exists(value["path"])
-    ):
-        value = PILImage.open(value["path"])
+    elif isinstance(value, dict) and "path" in value and isinstance(value["path"], str):
+        if os.path.exists(value["path"]):
+            value = PILImage.open(value["path"])
+        elif value["path"].startswith(f"hf://datasets/{dataset}@"):
+            with HfFileSystem(endpoint=hf_endpoint, token=hf_token).open(value["path"], "rb") as f:
+                src = value["path"].replace("hf://", hf_endpoint + "/", 1).replace("@", "/resolve/", 1)
+                image = PILImage.open(f)
+                return ImageSource(src=src, height=image.height, width=image.width)
+
     if not isinstance(value, PILImage.Image):
         raise TypeError(
             "Image cell must be a PIL image or an encoded dict of an image, "
@@ -141,6 +150,7 @@ def audio(
     value: Any,
     featureName: str,
     storage_client: StorageClient,
+    hf_endpoint: str,
     json_path: Optional[list[Union[str, int]]] = None,
 ) -> Any:
     from datasets.features._torchcodec import AudioDecoder
@@ -161,7 +171,14 @@ def audio(
             "Audio cell must be an encoded dict of an audio sample or a torchcodec AudioDecoder, "
             f"but got {str(value)[:300]}{'...' if len(str(value)) > 300 else ''}"
         )
+
     audio_file_extension = get_audio_file_extension(value)
+    if "path" in value and isinstance(value["path"], str) and value.get("bytes") is None:
+        if audio_file_extension in SUPPORTED_AUDIO_EXTENSION_TO_MEDIA_TYPE:
+            if value["path"].startswith(f"hf://datasets/{dataset}@"):
+                src = value["path"].replace("hf://", hf_endpoint + "/", 1).replace("@", "/resolve/", 1)
+                return AudioSource(src=src, type=SUPPORTED_AUDIO_EXTENSION_TO_MEDIA_TYPE[audio_file_extension])
+
     audio_file_bytes = get_audio_file_bytes(value)
     if not audio_file_extension:
         audio_file_extension = infer_audio_file_extension(audio_file_bytes)
@@ -262,6 +279,7 @@ def video(
     value: Any,
     featureName: str,
     storage_client: StorageClient,
+    hf_endpoint: str,
     json_path: Optional[list[Union[str, int]]] = None,
 ) -> Any:
     if datasets.config.TORCHCODEC_AVAILABLE:
@@ -292,6 +310,11 @@ def video(
             f"but got {str(value)[:300]}{'...' if len(str(value)) > 300 else ''}"
         )
 
+    if "path" in value and isinstance(value["path"], str) and value.get("bytes") is None:
+        if value["path"].startswith(f"hf://datasets/{dataset}@"):
+            src = value["path"].replace("hf://", hf_endpoint + "/", 1).replace("@", "/resolve/", 1)
+            return VideoSource(src=src)
+
     video_file_extension = get_video_file_extension(value)
     video_file_bytes = get_video_file_bytes(value)
     if not video_file_extension:
@@ -346,6 +369,8 @@ def pdf(
     value: Any,
     featureName: str,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     json_path: Optional[list[Union[str, int]]] = None,
 ) -> Any:
     if value is None:
@@ -354,13 +379,12 @@ def pdf(
         value = pdfplumber.open(BytesIO(value["bytes"]))
     elif isinstance(value, bytes):
         value = pdfplumber.open(BytesIO(value))
-    elif (
-        isinstance(value, dict)
-        and "path" in value
-        and isinstance(value["path"], str)
-        and os.path.exists(value["path"])
-    ):
-        value = pdfplumber.open(value["path"])
+    elif isinstance(value, dict) and "path" in value and isinstance(value["path"], str):
+        if os.path.exists(value["path"]):
+            value = pdfplumber.open(value["path"])
+        elif value["path"].startswith(f"hf://datasets/{dataset}@"):
+            f = HfFileSystem(endpoint=hf_endpoint, token=hf_token).open(value["path"], "rb")
+            value = pdfplumber.open(f)
 
     if not isinstance(value, pdfplumber.pdf.PDF):
         raise TypeError(
@@ -392,6 +416,8 @@ def get_cell_value(
     featureName: str,
     fieldType: Any,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     json_path: Optional[list[Union[str, int]]] = None,
 ) -> Any:
     # always allow None values in the cells
@@ -407,6 +433,8 @@ def get_cell_value(
             value=cell,
             featureName=featureName,
             storage_client=storage_client,
+            hf_endpoint=hf_endpoint,
+            hf_token=hf_token,
             json_path=json_path,
         )
     elif isinstance(fieldType, Audio):
@@ -419,6 +447,7 @@ def get_cell_value(
             value=cell,
             featureName=featureName,
             storage_client=storage_client,
+            hf_endpoint=hf_endpoint,
             json_path=json_path,
         )
     elif isinstance(fieldType, Video):
@@ -431,6 +460,7 @@ def get_cell_value(
             value=cell,
             featureName=featureName,
             storage_client=storage_client,
+            hf_endpoint=hf_endpoint,
             json_path=json_path,
         )
     elif isinstance(fieldType, Pdf):
@@ -443,6 +473,8 @@ def get_cell_value(
             value=cell,
             featureName=featureName,
             storage_client=storage_client,
+            hf_endpoint=hf_endpoint,
+            hf_token=hf_token,
             json_path=json_path,
         )
     elif isinstance(fieldType, Json):
@@ -467,6 +499,8 @@ def get_cell_value(
                 featureName=featureName,
                 fieldType=subFieldType,
                 storage_client=storage_client,
+                hf_endpoint=hf_endpoint,
+                hf_token=hf_token,
                 json_path=json_path + [idx] if json_path else [idx],
             )
             for (idx, subCell) in enumerate(cell)
@@ -486,6 +520,8 @@ def get_cell_value(
                 featureName=featureName,
                 fieldType=subFieldType,
                 storage_client=storage_client,
+                hf_endpoint=hf_endpoint,
+                hf_token=hf_token,
                 json_path=json_path + [idx] if json_path else [idx],
             )
             for (idx, subCell) in enumerate(cell)
@@ -506,6 +542,8 @@ def get_cell_value(
                 featureName=featureName,
                 fieldType=fieldType.feature,
                 storage_client=storage_client,
+                hf_endpoint=hf_endpoint,
+                hf_token=hf_token,
                 json_path=json_path + [idx] if json_path else [idx],
             )
             for (idx, subCell) in enumerate(cell)
@@ -525,6 +563,8 @@ def get_cell_value(
                 featureName=featureName,
                 fieldType=fieldType[key],
                 storage_client=storage_client,
+                hf_endpoint=hf_endpoint,
+                hf_token=hf_token,
                 json_path=json_path + [key] if json_path else [key],
             )
             for (key, subCell) in cell.items()

libs/libcommon/src/libcommon/viewer_utils/rows.py

@@ -2,7 +2,7 @@
 # Copyright 2022 The HuggingFace Authors.
 
 
-from typing import Protocol
+from typing import Optional, Protocol
 
 import PIL
 from datasets import Audio, Features, Image, Pdf, Value, Video
@@ -29,6 +29,8 @@ def transform_rows(
     rows: list[Row],
     features: Features,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
 ) -> list[Row]:
     transformed_rows: list[Row] = []
     for row_idx, row in enumerate(rows):
@@ -45,6 +47,8 @@ def transform_rows(
                         featureName=featureName,
                         fieldType=fieldType,
                         storage_client=storage_client,
+                        hf_endpoint=hf_endpoint,
+                        hf_token=hf_token,
                     )
                     for (featureName, fieldType) in features.items()
                 }
@@ -71,6 +75,8 @@ def create_first_rows_response(
     config: str,
     split: str,
     storage_client: StorageClient,
+    hf_endpoint: str,
+    hf_token: Optional[str],
     features: Features,
     get_rows_content: GetRowsContent,
     min_cell_bytes: int,
@@ -161,6 +167,8 @@ def create_first_rows_response(
         rows=rows_content.rows,
         features=features,
         storage_client=storage_client,
+        hf_endpoint=hf_endpoint,
+        hf_token=hf_token,
     )
 
     # truncate the rows to fit within the restrictions, and prepare them as RowItems