kedro-datasets

from copy import deepcopy
from pathlib import Path, PurePosixPath
from typing import Any, Dict
import fsspec
import numpy as np
from datasets import load_dataset
from datasets.arrow_dataset import Dataset
from <http://kedro.io|kedro.io> import AbstractDataSet
from kedro.io.core import get_filepath_str, get_protocol_and_path, DataSetError


class AudioFolderDataSet(AbstractDataSet[Dict[str, Any], Dataset]):
    """``AudioFolderDataSet`` loads audio data from the Hugging Face AudioFolder dataset.
    <https://huggingface.co/docs/datasets/audio_dataset#audiofolder>

    Example:
    ::

        >>> AudioFolderDataSet(data_dir='/path/to/data')
    """

    DEFAULT_LOAD_ARGS: Dict[str, Any] = {}
    DEFAULT_SAVE_ARGS: Dict[str, Any] = {}

    def _init_(
        self,
        filepath: str,
        load_args: Dict[str, Any] = None,
        save_args: Dict[str, Any] = None,
        credentials: Dict[str, Any] = None,
        fs_args: Dict[str, Any] = None,
        metadata: Dict[str, Any] = None,
    ):
        """Creates a new instance of AudioFolderDataSet to load audio data from the Hugging Face AudioFolder dataset.

        Args:
            data_dir: The location of the AudioFolder dataset.
        """
        protocol, self.path = get_protocol_and_path(filepath)
        self._protocol = protocol
        _fs_args = deepcopy(fs_args) or {}
        _fs_open_args_load = _fs_args.pop("open_args_load", {})
        _fs_open_args_save = _fs_args.pop("open_args_save", {})
        _credentials = deepcopy(credentials) or {}

        if protocol == "file":
            _fs_args.setdefault("auto_mkdir", True)

        self._fs = fsspec.filesystem(self._protocol, **_credentials, **_fs_args)

        self.metadata = metadata
        super()._init_()

        # super()._init_(
        #     filepath=PurePosixPath(path),
        #     exists_function=self._fs.exists,
        #     glob_function=self._fs.glob,
        # )

        # Handle default load and save arguments
        self._load_args = deepcopy(self.DEFAULT_LOAD_ARGS)
        if load_args is not None:
            self._load_args.update(load_args)
        self._save_args = deepcopy(self.DEFAULT_SAVE_ARGS)
        if save_args is not None:
            self._save_args.update(save_args)

        _fs_open_args_save.setdefault("mode", "wb")
        self._fs_open_args_load = _fs_open_args_load
        self._fs_open_args_save = _fs_open_args_save

        self._filepath = get_filepath_str(PurePosixPath(self.path), self._protocol)

    def _load(self) -> Dict[str, Any]:
        """Loads data from the AudioFolder dataset.

        Returns:
            Data from the AudioFolder dataset as a dictionary of train, validation, and test sets.
        """

        # with self._fs.open(load_path, **self._fs_open_args_load) as fs_file:
        #     fs_path = fs_file.path
        return load_dataset("audiofolder", data_dir=self._filepath)

    def _save(self, data: Dict[str, Any]) -> None:
        """Saves audio data to the specified filepath."""
        raise NotImplementedError("AudioFolderDataSet does not support saving data.")

    def _describe(self) -> Dict[str, Any]:
        """Returns a dict that describes the attributes of the dataset."""
        return dict(filepath=self._filepath)

    def _exists(self) -> bool:
        return self._fs.exists(self._filepath)

from tempfile import TemporaryDirectory
from cloudpathlib import CloudPath
from transformers import AutoModel
from <http://kedro.io|kedro.io> import AbstractDataSet

class HFTransformersDataset(AbstractDataSet):
    def __init__(self, path: str):
        self.cloud_path = CloudPath(path)

    def _load(self) -> AutoModel:
        """
        Loads the model from the cloud path.

        Returns:
            AutoModel: The loaded Hugging Face Transformers model.
        """
        with TemporaryDirectory() as tmp_dir:
            self.cloud_path.download_to(tmp_dir)
            return AutoModel.from_pretrained(tmp_dir)

    def _save(self, model: AutoModel) -> None:
        """
        Saves the model to the cloud path.

        Args:
            model (AutoModel): The Hugging Face Transformers model to save.
        """
        with TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir)
            self.cloud_path.upload_from(tmp_dir)

    def _describe(self) -> Dict[str, Any]:
        return dict(
            cloud_path=str(self.cloud_path),
        )