dataeval 0.72.0__py3-none-any.whl → 0.72.2__py3-none-any.whl

dataeval/{_internal/models/pytorch → utils/torch}/blocks.py

@@ -1,3 +1,7 @@
+from typing import Any
+
+__all__ = []
+
 import torch.nn as nn
 
 
@@ -8,21 +12,22 @@ class Conv(nn.Module):
 
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        k=1,
-        s=1,
-        p=0,
-        activation="relu",
-        norm="instance",
-    ):
+        in_channels: int,
+        out_channels: int,
+        k: int = 1,
+        s: int = 1,
+        p: int = 0,
+        activation: str = "relu",
+        norm: str = "instance",
+    ) -> None:
         super().__init__()
-        conv = nn.Conv2d(in_channels, out_channels, kernel_size=k, stride=s, padding=p)
-        norm = self.get_norm_func(norm=norm, out_channels=out_channels)
-        act = self.get_activation_func(activation=activation)
-        self.module = nn.Sequential(conv, norm, act)
+        self.module: nn.Sequential = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=k, stride=s, padding=p),
+            self.get_norm_func(norm=norm, out_channels=out_channels),
+            self.get_activation_func(activation=activation),
+        )
 
-    def get_norm_func(self, norm: str, out_channels) -> nn.Module:
+    def get_norm_func(self, norm: str, out_channels: int) -> nn.Module:
         if norm == "batch":
             return nn.BatchNorm2d(out_channels)
         if norm == "instance":
@@ -42,5 +47,5 @@ class Conv(nn.Module):
             return nn.Tanh()
         return nn.Identity()
 
-    def forward(self, x):
+    def forward(self, x: Any) -> Any:
         return self.module(x)

dataeval/{_internal → utils/torch}/datasets.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+__all__ = ["MNIST", "CIFAR10", "VOCDetection"]
+
 import hashlib
 import os
 import zipfile
@@ -11,7 +13,7 @@ import numpy as np
 import requests
 from numpy.typing import NDArray
 from torch.utils.data import Dataset
-from torchvision.datasets import CIFAR10, VOCDetection  # noqa: F401
+from torchvision.datasets import CIFAR10, VOCDetection
 
 ClassStringMap = Literal["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 TClassMap = TypeVar("TClassMap", ClassStringMap, int, list[ClassStringMap], list[int])
@@ -50,6 +52,7 @@ def _get_file(
     file_hash: str | None = None,
     verbose: bool = True,
     md5: bool = False,
+    timeout: int = 60,
 ):
     fpath = os.path.join(root, fname)
     download = True
@@ -64,16 +67,16 @@ def _get_file(
         try:
             error_msg = "URL fetch failure on {}: {} -- {}"
             try:
-                with requests.get(origin, stream=True, timeout=60) as r:
+                with requests.get(origin, stream=True, timeout=timeout) as r:
                     r.raise_for_status()
                     with open(fpath, "wb") as f:
                         for chunk in r.iter_content(chunk_size=8192):
                             if chunk:
                                 f.write(chunk)
             except requests.exceptions.HTTPError as e:
-                raise Exception(f"{error_msg.format(origin, e.response.status_code)} -- {e.response.reason}") from e
+                raise RuntimeError(f"{error_msg.format(origin, e.response.status_code, e.response.reason)}") from e
             except requests.exceptions.RequestException as e:
-                raise Exception(f"{error_msg.format(origin, 'Unknown error')} -- {str(e)}") from e
+                raise ValueError(f"{error_msg.format(origin, 'Unknown error', str(e))}") from e
         except (Exception, KeyboardInterrupt):
             if os.path.exists(fpath):
                 os.remove(fpath)
@@ -89,7 +92,7 @@ def _get_file(
     return fpath
 
 
-def check_exists(
+def _check_exists(
     folder: str | Path,
     url: str,
     root: str | Path,
@@ -103,7 +106,7 @@ def check_exists(
     location = str(folder)
     if not os.path.exists(folder):
         if download:
-            location = download_dataset(url, root, fname, file_hash, verbose, md5)
+            location = _download_dataset(url, root, fname, file_hash, verbose, md5)
         else:
             raise RuntimeError("Dataset not found. You can use download=True to download it")
     else:
@@ -112,7 +115,7 @@ def check_exists(
     return location
 
 
-def download_dataset(
+def _download_dataset(
     url: str, root: str | Path, fname: str, file_hash: str, verbose: bool = True, md5: bool = False
 ) -> str:
     """Code to download mnist and corruptions, originates from tensorflow_datasets (tfds):
@@ -131,11 +134,11 @@ def download_dataset(
         md5=md5,
     )
     if md5:
-        folder = extract_archive(fpath, root, remove_finished=True)
+        folder = _extract_archive(fpath, root, remove_finished=True)
     return folder
 
 
-def extract_archive(
+def _extract_archive(
     from_path: str | Path,
     to_path: str | Path | None = None,
     remove_finished: bool = False,
@@ -163,13 +166,13 @@ def extract_archive(
     return str(to_path)
 
 
-def subselect(arr: NDArray, count: int, from_back: bool = False):
+def _subselect(arr: NDArray, count: int, from_back: bool = False):
     if from_back:
         return arr[-count:]
     return arr[:count]
 
 
-class MNIST(Dataset):
+class MNIST(Dataset[tuple[NDArray[np.float64], int]]):
     """MNIST Dataset and Corruptions.
 
     Args:
@@ -185,7 +188,7 @@ class MNIST(Dataset):
         unit_interval : bool, default False
             Shift the data values to the unit interval [0-1].
         dtype : type | None, default None
-            Change the numpy dtype - data is loaded as np.uint8
+            Change the :term:`NumPy` dtype - data is loaded as np.uint8
         channels : Literal['channels_first' | 'channels_last'] | None, default None
             Location of channel axis if desired, default has no channels (N, 28, 28)
         flatten : bool, default False
@@ -211,17 +214,17 @@ class MNIST(Dataset):
             If True, outputs print statements.
     """
 
-    mirror = [
+    _mirrors: tuple[str, ...] = (
         "https://storage.googleapis.com/tensorflow/tf-keras-datasets/",
         "https://zenodo.org/record/3239543/files/",
-    ]
+    )
 
-    resources = [
+    _resources: tuple[tuple[str, str], ...] = (
         ("mnist.npz", "731c5ac602752760c8e48fbffcf8c3b850d9dc2a2aedcf2cc48468fc17b673d1"),
         ("mnist_c.zip", "4b34b33045869ee6d424616cd3a65da3"),
-    ]
+    )
 
-    class_dict = {
+    class_dict: dict[str, int] = {
         "zero": 0,
         "one": 1,
         "two": 2,
@@ -267,43 +270,46 @@ class MNIST(Dataset):
         self.randomize = randomize
         self.from_back = slice_back
         self.verbose = verbose
+        self.data: NDArray[np.float64]
+        self.targets: NDArray[np.int_]
+        self.size: int
 
-        self.class_set = []
+        self._class_set = []
         if classes is not None:
             if not isinstance(classes, list):
                 classes = [classes]  # type: ignore
 
             for val in classes:  # type: ignore
                 if isinstance(val, int) and 0 <= val < 10:
-                    self.class_set.append(val)
+                    self._class_set.append(val)
                 elif isinstance(val, str):
-                    self.class_set.append(self.class_dict[val])
-            self.class_set = set(self.class_set)
+                    self._class_set.append(self.class_dict[val])
+            self._class_set = set(self._class_set)
 
-        if not self.class_set:
-            self.class_set = set(self.class_dict.values())
+        if not self._class_set:
+            self._class_set = set(self.class_dict.values())
 
-        self.num_classes = len(self.class_set)
+        self._num_classes = len(self._class_set)
 
         if self.corruption is None:
-            file_resource = self.resources[0]
-            mirror = self.mirror[0]
+            file_resource = self._resources[0]
+            mirror = self._mirrors[0]
             md5 = False
         else:
             if self.corruption == "identity" and verbose:
                 print("Identity is not a corrupted dataset but the original MNIST dataset.")
-            file_resource = self.resources[1]
-            mirror = self.mirror[1]
+            file_resource = self._resources[1]
+            mirror = self._mirrors[1]
             md5 = True
-        check_exists(self.mnist_folder, mirror, self.root, file_resource[0], file_resource[1], download, verbose, md5)
+        _check_exists(self.mnist_folder, mirror, self.root, file_resource[0], file_resource[1], download, verbose, md5)
 
         self.data, self.targets = self._load_data()
 
         self._augmentations()
 
-    def _load_data(self):
+    def _load_data(self) -> tuple[NDArray[np.float64], NDArray[np.int64]]:
         if self.corruption is None:
-            image_file = self.resources[0][0]
+            image_file = self._resources[0][0]
             data, targets = self._read_normal_file(os.path.join(self.mnist_folder, image_file))
         else:
             image_file = f"{'train' if self.train else 'test'}_images.npy"
@@ -329,27 +335,27 @@ class MNIST(Dataset):
             self.data = self.data[shuffled_indices]
             self.targets = self.targets[shuffled_indices]
 
-        if not self.balance and self.num_classes > self.size:
+        if not self.balance and self._num_classes > self.size:
             if self.size > 0:
-                self.data = subselect(self.data, self.size, self.from_back)
-                self.targets = subselect(self.targets, self.size, self.from_back)
+                self.data = _subselect(self.data, self.size, self.from_back)
+                self.targets = _subselect(self.targets, self.size, self.from_back)
         else:
-            label_dict = {label: np.where(self.targets == label)[0] for label in self.class_set}
+            label_dict = {label: np.where(self.targets == label)[0] for label in self._class_set}
             min_label_count = min(len(indices) for indices in label_dict.values())
 
-            self.per_class_count = int(np.ceil(self.size / self.num_classes)) if self.size > 0 else min_label_count
+            self._per_class_count = int(np.ceil(self.size / self._num_classes)) if self.size > 0 else min_label_count
 
-            if self.per_class_count > min_label_count:
-                self.per_class_count = min_label_count
+            if self._per_class_count > min_label_count:
+                self._per_class_count = min_label_count
                 if not self.balance and self.verbose:
                     warn(
-                        f"Because of dataset limitations, only {min_label_count*self.num_classes} samples "
+                        f"Because of dataset limitations, only {min_label_count*self._num_classes} samples "
                         f"will be returned, instead of the desired {self.size}."
                     )
 
-            all_indices = np.empty(shape=(self.num_classes, self.per_class_count), dtype=int)
-            for i, label in enumerate(self.class_set):
-                all_indices[i] = subselect(label_dict[label], self.per_class_count, self.from_back)
+            all_indices: NDArray[np.int_] = np.empty(shape=(self._num_classes, self._per_class_count), dtype=np.int_)
+            for i, label in enumerate(self._class_set):
+                all_indices[i] = _subselect(label_dict[label], self._per_class_count, self.from_back)
             self.data = np.vstack(self.data[all_indices.T])  # type: ignore
             self.targets = np.hstack(self.targets[all_indices.T])  # type: ignore
 
@@ -370,7 +376,7 @@ class MNIST(Dataset):
         if self.flatten and self.channels is None:
             self.data = self.data.reshape(self.data.shape[0], -1)
 
-    def __getitem__(self, index: int) -> tuple[NDArray, int]:
+    def __getitem__(self, index: int) -> tuple[NDArray[np.float64], int]:
         """
         Args:
             index (int): Index

dataeval/utils/torch/models.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+__all__ = ["AriaAutoencoder", "Encoder", "Decoder"]
+
+from typing import Any
+
+import torch.nn as nn
+
+
+class AriaAutoencoder(nn.Module):
+    """
+    An autoencoder model with a separate encoder and decoder.
+
+    Parameters
+    ----------
+    channels : int, default 3
+        Number of input channels
+    """
+
+    def __init__(self, channels: int = 3) -> None:
+        super().__init__()
+        self.encoder: Encoder = Encoder(channels)
+        self.decoder: Decoder = Decoder(channels)
+
+    def forward(self, x: Any) -> Any:
+        """
+        Perform a forward pass through the encoder and decoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor
+
+        Returns
+        -------
+        torch.Tensor
+            The reconstructed output tensor.
+        """
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+    def encode(self, x: Any) -> Any:
+        """
+        Encode the input tensor using the encoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor
+
+        Returns
+        -------
+        torch.Tensor
+            The encoded representation of the input tensor.
+        """
+        return self.encoder(x)
+
+
+class Encoder(nn.Module):
+    """
+    A simple encoder to be used in an autoencoder model.
+
+    This is the encoder used by the AriaAutoencoder model.
+
+    Parameters
+    ----------
+    channels : int, default 3
+        Number of input channels
+    """
+
+    def __init__(self, channels: int = 3) -> None:
+        super().__init__()
+        self.encoder: nn.Sequential = nn.Sequential(
+            nn.Conv2d(channels, 256, 2, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(256, 128, 2, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(128, 64, 2, stride=1),
+        )
+
+    def forward(self, x: Any) -> Any:
+        """
+        Perform a forward pass through the encoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor
+
+        Returns
+        -------
+        torch.Tensor
+            The encoded representation of the input tensor.
+        """
+        return self.encoder(x)
+
+
+class Decoder(nn.Module):
+    """
+    A simple decoder to be used in an autoencoder model.
+
+    This is the decoder used by the AriaAutoencoder model.
+
+    Parameters
+    ----------
+    channels : int
+        Number of output channels
+    """
+
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.decoder: nn.Sequential = nn.Sequential(
+            nn.ConvTranspose2d(64, 128, 2, stride=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(128, 256, 2, stride=2),
+            nn.ReLU(),
+            nn.ConvTranspose2d(256, channels, 2, stride=2),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: Any) -> Any:
+        """
+        Perform a forward pass through the decoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The encoded tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            The reconstructed output tensor.
+        """
+        return self.decoder(x)

dataeval/{_internal/models/pytorch/autoencoder.py → utils/torch/trainer.py}

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+__all__ = ["AETrainer"]
+
 from typing import Any
 
 import torch
@@ -17,7 +19,7 @@ def get_images_from_batch(batch: Any) -> Any:
 
 class AETrainer:
     """
-    A class to train and evaluate an autoencoder model.
+    A class to train and evaluate an autoencoder<Autoencoder>` model.
 
     Parameters
     ----------
@@ -38,13 +40,13 @@ class AETrainer:
     ):
         if device == "auto":
             device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.device = device
-        self.model = model.to(device)
+        self.device: torch.device = torch.device(device)
+        self.model: nn.Module = model.to(device)
         self.batch_size = batch_size
 
-    def train(self, dataset: Dataset, epochs: int = 25) -> list[float]:
+    def train(self, dataset: Dataset[Any], epochs: int = 25) -> list[float]:
         """
-        Basic image reconstruction training function for Autoencoder models
+        Basic image reconstruction training function for :term:`Autoencoder` models
 
         Uses `torch.optim.Adam` and `torch.nn.MSELoss` as default hyperparameters
 
@@ -59,7 +61,7 @@ class AETrainer:
         Returns
         -------
         List[float]
-            A list of average loss values for each epoch.
+            A list of average loss values for each :term:`epoch<Epoch>`.
 
         Note
         ----
@@ -101,9 +103,9 @@ class AETrainer:
         return loss_history
 
     @torch.no_grad
-    def eval(self, dataset: Dataset) -> float:
+    def eval(self, dataset: Dataset[Any]) -> float:
         """
-        Basic image reconstruction evaluation function for Autoencoder models
+        Basic image reconstruction evaluation function for :term:`autoencoder<Autoencoder>` models
 
         Uses `torch.nn.MSELoss` as default loss function.
 
@@ -137,9 +139,9 @@ class AETrainer:
         return total_loss / len(dataloader)
 
     @torch.no_grad
-    def encode(self, dataset: Dataset) -> torch.Tensor:
+    def encode(self, dataset: Dataset[Any]) -> torch.Tensor:
         """
-        Create image embeddings for the dataset using the model's encoder.
+        Create image :term:`embeddings<Embeddings>` for the dataset using the model's encoder.
 
         If the model has an `encode` method, it will be used; otherwise,
         `model.forward` will be used.
@@ -174,134 +176,3 @@ class AETrainer:
             encodings = torch.vstack((encodings, embeddings)) if len(encodings) else embeddings
 
         return encodings
-
-
-class AriaAutoencoder(nn.Module):
-    """
-    An autoencoder model with a separate encoder and decoder.
-
-    Parameters
-    ----------
-    channels : int, default 3
-        Number of input channels
-    """
-
-    def __init__(self, channels=3):
-        super().__init__()
-        self.encoder = Encoder(channels)
-        self.decoder = Decoder(channels)
-
-    def forward(self, x):
-        """
-        Perform a forward pass through the encoder and decoder.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input tensor
-
-        Returns
-        -------
-        torch.Tensor
-            The reconstructed output tensor.
-        """
-        x = self.encoder(x)
-        x = self.decoder(x)
-        return x
-
-    def encode(self, x):
-        """
-        Encode the input tensor using the encoder.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input tensor
-
-        Returns
-        -------
-        torch.Tensor
-            The encoded representation of the input tensor.
-        """
-        return self.encoder(x)
-
-
-class Encoder(nn.Module):
-    """
-    A simple encoder to be used in an autoencoder model.
-
-    This is the encoder used by the AriaAutoencoder model.
-
-    Parameters
-    ----------
-    channels : int, default 3
-        Number of input channels
-    """
-
-    def __init__(self, channels=3):
-        super().__init__()
-        self.encoder = nn.Sequential(
-            nn.Conv2d(channels, 256, 2, stride=1, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Conv2d(256, 128, 2, stride=1, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Conv2d(128, 64, 2, stride=1),
-        )
-
-    def forward(self, x):
-        """
-        Perform a forward pass through the encoder.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input tensor
-
-        Returns
-        -------
-        torch.Tensor
-            The encoded representation of the input tensor.
-        """
-        return self.encoder(x)
-
-
-class Decoder(nn.Module):
-    """
-    A simple decoder to be used in an autoencoder model.
-
-    This is the decoder used by the AriaAutoencoder model.
-
-    Parameters
-    ----------
-    channels : int
-        Number of output channels
-    """
-
-    def __init__(self, channels):
-        super().__init__()
-        self.decoder = nn.Sequential(
-            nn.ConvTranspose2d(64, 128, 2, stride=1),
-            nn.ReLU(),
-            nn.ConvTranspose2d(128, 256, 2, stride=2),
-            nn.ReLU(),
-            nn.ConvTranspose2d(256, channels, 2, stride=2),
-            nn.Sigmoid(),
-        )
-
-    def forward(self, x):
-        """
-        Perform a forward pass through the decoder.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            The encoded tensor.
-
-        Returns
-        -------
-        torch.Tensor
-            The reconstructed output tensor.
-        """
-        return self.decoder(x)

dataeval/{_internal → utils/torch}/utils.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+__all__ = ["read_dataset"]
+
 from collections import defaultdict
 from typing import Any
 
 from torch.utils.data import Dataset
 
 
-def read_dataset(dataset: Dataset) -> list[list[Any]]:
+def read_dataset(dataset: Dataset[Any]) -> list[list[Any]]:
     """
     Extract information from a dataset at each index into individual lists of each information position
 

dataeval/workflows/__init__.py

@@ -5,6 +5,6 @@ Workflows perform a sequence of actions to analyze the dataset and make predicti
 from dataeval import _IS_TORCH_AVAILABLE
 
 if _IS_TORCH_AVAILABLE:  # pragma: no cover
-    from dataeval._internal.workflows.sufficiency import Sufficiency, SufficiencyOutput
+    from dataeval.workflows.sufficiency import Sufficiency, SufficiencyOutput
 
     __all__ = ["Sufficiency", "SufficiencyOutput"]