Mod: Update transform typing with item typeddict, dataloader example …

…in notebook and version in init.

examples/dataloader.ipynb

@@ -60,7 +60,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "aac-datasets version: 0.5.0\n"
+      "aac-datasets version: 0.5.1\n"
      ]
     }
    ],
@@ -108,7 +108,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -123,7 +123,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -188,7 +188,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.10.11"
   },
   "orig_nbformat": 4,
   "vscode": {

src/aac_datasets/__init__.py

@@ -12,7 +12,7 @@
 __maintainer__ = "Etienne Labbé (Labbeti)"
 __name__ = "aac-datasets"
 __status__ = "Development"
-__version__ = "0.5.0"
+__version__ = "0.5.1"
 
 
 from .datasets.audiocaps import AudioCaps
@@ -28,7 +28,6 @@
     set_default_ytdlp_path,
 )
 
-
 __all__ = [
     "AudioCaps",
     "Clotho",

src/aac_datasets/datasets/audiocaps.py

@@ -87,7 +87,7 @@ def __init__(
         root: Union[str, Path, None] = None,
         subset: str = AudioCapsCard.DEFAULT_SUBSET,
         download: bool = False,
-        transform: Optional[Callable[[Dict[str, Any]], Any]] = None,
+        transform: Optional[Callable[[AudioCapsItem], Any]] = None,
         verbose: int = 0,
         force_download: bool = False,
         verify_files: bool = False,

src/aac_datasets/datasets/base.py

@@ -19,7 +19,6 @@
 
 import torchaudio
 import tqdm
-from typing_extensions import TypedDict
 
 try:
     # To support torchaudio >= 2.1.0
@@ -33,7 +32,7 @@
 pylog = logging.getLogger(__name__)
 
 
-ItemType = TypeVar("ItemType", bound=TypedDict, covariant=True)
+ItemType = TypeVar("ItemType", covariant=True)
 
 
 class AACDataset(Generic[ItemType], Dataset[ItemType]):
@@ -43,7 +42,7 @@ class AACDataset(Generic[ItemType], Dataset[ItemType]):
     def __init__(
         self,
         raw_data: Optional[Dict[str, List[Any]]] = None,
-        transform: Optional[Callable] = None,
+        transform: Optional[Callable[[ItemType], Any]] = None,
         column_names: Optional[Iterable[str]] = None,
         flat_captions: bool = False,
         sr: Optional[int] = None,
@@ -146,7 +145,7 @@ def column_names(
         self._columns = columns
 
     @transform.setter
-    def transform(self, transform: Optional[Callable]) -> None:
+    def transform(self, transform: Optional[Callable[[ItemType], Any]]) -> None:
         self._transform = transform
 
     # Public methods
@@ -158,6 +157,10 @@ def at(self, index: int) -> ItemType:
     def at(self, index: Union[Iterable[int], slice, None], column: str) -> List:
         ...
 
+    @overload
+    def at(self, index: Union[Iterable[int], slice, None]) -> Dict[str, List]:
+        ...
+
     @overload
     def at(
         self,
@@ -376,10 +379,10 @@ def __getitem__(self, index: Any) -> Any:
         item = self.at(index, column)
         if (
             isinstance(index, int)
-            and (column is None or column == self._columns)
             and self._transform is not None
+            and (column is None or set(column) == set(self._columns))
         ):
-            item = self._transform(item)
+            item = self._transform(item)  # type: ignore
         return item
 
     def __len__(self) -> int:
@@ -394,7 +397,7 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         info = {
             "size": len(self),
-            "num_columns": len(self.column_names),
+            "num_columns": self.num_columns,
         }
         repr_str = ", ".join(f"{k}={v}" for k, v in info.items())
         return f"{self.__class__.__name__}({repr_str})"

src/aac_datasets/datasets/clotho.py

@@ -3,29 +3,21 @@
 
 import logging
 import os.path as osp
-
 from pathlib import Path
-from typing import (
-    Callable,
-    ClassVar,
-    List,
-    Optional,
-    Union,
-)
+from typing import Any, Callable, ClassVar, List, Optional, Union
 
 from torch import Tensor
-from typing_extensions import TypedDict, NotRequired
+from typing_extensions import NotRequired, TypedDict
 
 from aac_datasets.datasets.base import AACDataset
 from aac_datasets.datasets.functional.clotho import (
     ClothoCard,
-    load_clotho_dataset,
-    download_clotho_dataset,
     _get_audio_subset_dpath,
+    download_clotho_dataset,
+    load_clotho_dataset,
 )
 from aac_datasets.utils.globals import _get_root
 
-
 pylog = logging.getLogger(__name__)
 
 
@@ -108,7 +100,7 @@ def __init__(
         root: Union[str, Path, None] = None,
         subset: str = ClothoCard.DEFAULT_SUBSET,
         download: bool = False,
-        transform: Optional[Callable] = None,
+        transform: Optional[Callable[[ClothoItem], Any]] = None,
         verbose: int = 0,
         force_download: bool = False,
         verify_files: bool = False,

src/aac_datasets/datasets/macs.py

@@ -5,7 +5,7 @@
 import logging
 import os.path as osp
 from pathlib import Path
-from typing import Callable, ClassVar, Dict, List, Optional, Union
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Union
 
 from torch import Tensor
 from typing_extensions import TypedDict
@@ -71,7 +71,7 @@ def __init__(
         root: Union[str, Path, None] = None,
         subset: str = MACSCard.DEFAULT_SUBSET,
         download: bool = False,
-        transform: Optional[Callable] = None,
+        transform: Optional[Callable[[MACSItem], Any]] = None,
         verbose: int = 0,
         force_download: bool = False,
         verify_files: bool = False,

src/aac_datasets/datasets/wavcaps.py

@@ -3,23 +3,21 @@
 
 import logging
 import os.path as osp
-
 from pathlib import Path
-from typing import Callable, ClassVar, List, Optional, Union
+from typing import Any, Callable, ClassVar, List, Optional, Union
 
 from torch import Tensor
 from typing_extensions import TypedDict
 
 from aac_datasets.datasets.base import AACDataset
 from aac_datasets.datasets.functional.wavcaps import (
     WavCapsCard,
-    load_wavcaps_dataset,
-    download_wavcaps_dataset,
     _get_audio_subset_dpath,
+    download_wavcaps_dataset,
+    load_wavcaps_dataset,
 )
 from aac_datasets.utils.globals import _get_root, _get_zip_path
 
-
 pylog = logging.getLogger(__name__)
 
 
@@ -110,7 +108,7 @@ def __init__(
         root: Union[str, Path, None] = None,
         subset: str = WavCapsCard.DEFAULT_SUBSET,
         download: bool = False,
-        transform: Optional[Callable] = None,
+        transform: Optional[Callable[[WavCapsItem], Any]] = None,
         verbose: int = 0,
         force_download: bool = False,
         verify_files: bool = False,