dataeval 0.65.0__py3-none-any.whl → 0.67.0__py3-none-any.whl

dataeval/_internal/models/pytorch/autoencoder.py

@@ -1,4 +1,6 @@
-from typing import Any, List, Union
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -14,40 +16,52 @@ def get_images_from_batch(batch: Any) -> Any:
 
 
 class AETrainer:
+    """
+    A class to train and evaluate an autoencoder model.
+
+    Parameters
+    ----------
+    model : nn.Module
+        The model to be trained.
+    device : str or torch.device, default "auto"
+        The hardware device to use for training.
+        If "auto", the device will be set to "cuda" if available, otherwise "cpu".
+    batch_size : int, default 8
+        The number of images to process in a batch.
+    """
+
     def __init__(
         self,
         model: nn.Module,
-        device: Union[str, torch.device] = "auto",
+        device: str | torch.device = "auto",
         batch_size: int = 8,
     ):
-        """
-        model : nn.Module
-            Model to be trained
-        device : str | torch.device, default "cpu"
-            Hardware device for model, optimizer, and data to run on
-        batch_size : int, default 8
-            Number of images to group together in `torch.utils.data.DataLoader`
-        """
         if device == "auto":
             device = "cuda" if torch.cuda.is_available() else "cpu"
         self.device = device
         self.model = model.to(device)
         self.batch_size = batch_size
 
-    def train(self, dataset: Dataset, epochs: int = 25) -> List[float]:
+    def train(self, dataset: Dataset, epochs: int = 25) -> list[float]:
         """
-        Basic training function for Autoencoder models for reconstruction tasks
+        Basic image reconstruction training function for Autoencoder models
 
         Uses `torch.optim.Adam` and `torch.nn.MSELoss` as default hyperparameters
 
         Parameters
         ----------
         dataset : Dataset
-            Torch Dataset containing images in the first return position
+            The dataset to train on.
+            Torch Dataset containing images in the first return position.
         epochs : int, default 25
             Number of full training loops
 
-        Note
+        Returns
+        -------
+        List[float]
+            A list of average loss values for each epoch.
+
+        Notes
         ----
         To replace this function with a custom function, do
             AETrainer.train = custom_function
@@ -58,7 +72,7 @@ class AETrainer:
         opt = Adam(self.model.parameters(), lr=0.001)
         criterion = nn.MSELoss().to(self.device)
         # Record loss
-        loss_history: List[float] = []
+        loss_history: list[float] = []
 
         for _ in range(epochs):
             epoch_loss: float = 0
@@ -89,19 +103,20 @@ class AETrainer:
     @torch.no_grad
     def eval(self, dataset: Dataset) -> float:
         """
-        Basic evaluation function for Autoencoder models for reconstruction tasks
+        Basic image reconstruction evaluation function for Autoencoder models
 
-        Uses `torch.optim.Adam` and `torch.nn.MSELoss` as default hyperparameters
+        Uses `torch.nn.MSELoss` as default loss function.
 
         Parameters
         ----------
         dataset : Dataset
-            Torch Dataset containing images in the first return position
+            The dataset to evaluate on.
+            Torch Dataset containing images in the first return position.
 
         Returns
         -------
         float
-            Total reconstruction loss over all data
+            Total reconstruction loss over the entire dataset
 
         Note
         ----
@@ -124,18 +139,25 @@ class AETrainer:
     @torch.no_grad
     def encode(self, dataset: Dataset) -> torch.Tensor:
         """
-        Encode data through model if it has an encode attribute,
-        otherwise passes data through model.forward
+        Create image 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.
 
         Parameters
         ----------
         dataset: Dataset
-            Dataset containing images to be encoded by the model
+            The dataset to encode.
+            Torch Dataset containing images in the first return position.
 
         Returns
         -------
         torch.Tensor
             Data encoded by the model
+
+        Notes
+        -----
+        This function should be run after the model has been trained and evaluated.
         """
         self.model.eval()
         dl = DataLoader(dataset, batch_size=self.batch_size)
@@ -155,21 +177,67 @@ class AETrainer:
 
 
 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(
@@ -183,10 +251,34 @@ class Encoder(nn.Module):
         )
 
     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(
@@ -199,4 +291,17 @@ class Decoder(nn.Module):
         )
 
     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/models/tensorflow/autoencoder.py

@@ -8,7 +8,9 @@ Licensed under Apache Software License (Apache 2.0)
 
 # pyright: reportIncompatibleMethodOverride=false
 
-from typing import Callable, Tuple, cast
+from __future__ import annotations
+
+from typing import Callable, cast
 
 import keras
 import tensorflow as tf
@@ -56,16 +58,17 @@ def eucl_cosim_features(x: tf.Tensor, y: tf.Tensor, max_eucl: float = 1e2) -> tf
 
     Parameters
     ----------
-    x
+    x : tf.Tensor
         Tensor used in feature computation.
-    y
+    y : tf.Tensor
         Tensor used in feature computation.
-    max_eucl
+    max_eucl : float, default 1e2
         Maximum value to clip relative Euclidean distance by.
 
     Returns
     -------
-    Tensor concatenating the relative Euclidean distance and cosine similarity features.
+    tf.Tensor
+        Tensor concatenating the relative Euclidean distance and cosine similarity features.
     """
     if len(x.shape) > 2 or len(y.shape) > 2:
         x = cast(tf.Tensor, Flatten()(x))
@@ -78,9 +81,9 @@ def eucl_cosim_features(x: tf.Tensor, y: tf.Tensor, max_eucl: float = 1e2) -> tf
 
 
 class Sampling(Layer):
-    """Reparametrization trick. Uses (z_mean, z_log_var) to sample the latent vector z."""
+    """Reparametrization trick - Uses (z_mean, z_log_var) to sample the latent vector z."""
 
-    def call(self, inputs: Tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
+    def call(self, inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
         """
         Sample z.
 
@@ -138,7 +141,7 @@ class EncoderVAE(Layer):
         self.fc_log_var = Dense(latent_dim, activation=None)
         self.sampling = Sampling()
 
-    def call(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+    def call(self, x: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
         x = cast(tf.Tensor, self.encoder_net(x))
         if len(x.shape) > 2:
             x = cast(tf.Tensor, Flatten()(x))
@@ -173,9 +176,9 @@ class AE(keras.Model):
 
     Parameters
     ----------
-    encoder_net
+    encoder_net : keras.Model
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net
+    decoder_net : keras.Model
         Layers for the decoder wrapped in a keras.Sequential class.
     """
 
@@ -196,13 +199,13 @@ class VAE(keras.Model):
 
     Parameters
     ----------
-    encoder_net
+    encoder_net : keras.Model
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net
+    decoder_net : keras.Model
         Layers for the decoder wrapped in a keras.Sequential class.
-    latent_dim
+    latent_dim : int
         Dimensionality of the latent space.
-    beta
+    beta : float, default 1.0
         Beta parameter for KL-divergence loss term.
     """
 
@@ -214,7 +217,7 @@ class VAE(keras.Model):
         self.latent_dim = latent_dim
 
     def call(self, x: tf.Tensor) -> tf.Tensor:
-        z_mean, z_log_var, z = cast(Tuple[tf.Tensor, tf.Tensor, tf.Tensor], self.encoder(x))
+        z_mean, z_log_var, z = cast(tuple[tf.Tensor, tf.Tensor, tf.Tensor], self.encoder(x))
         x_recon = self.decoder(z)
         # add KL divergence loss term
         kl_loss = -0.5 * tf.reduce_mean(z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
@@ -228,15 +231,15 @@ class AEGMM(keras.Model):
 
     Parameters
     ----------
-    encoder_net
+    encoder_net : keras.Model
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net
+    decoder_net : keras.Model
         Layers for the decoder wrapped in a keras.Sequential class.
-    gmm_density_net
+    gmm_density_net : keras.Model
         Layers for the GMM network wrapped in a keras.Sequential class.
-    n_gmm
+    n_gmm : int
         Number of components in GMM.
-    recon_features
+    recon_features : Callable, default eucl_cosim_features
         Function to extract features from the reconstructed instance by the decoder.
     """
 
@@ -255,7 +258,7 @@ class AEGMM(keras.Model):
         self.n_gmm = n_gmm
         self.recon_features = recon_features
 
-    def call(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+    def call(self, x: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
         enc = self.encoder(x)
         x_recon = cast(tf.Tensor, self.decoder(enc))
         recon_features = self.recon_features(x, x_recon)
@@ -270,19 +273,19 @@ class VAEGMM(keras.Model):
 
     Parameters
     ----------
-    encoder_net
+    encoder_net : keras.Model
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net
+    decoder_net : keras.Model
         Layers for the decoder wrapped in a keras.Sequential class.
-    gmm_density_net
+    gmm_density_net : keras.Model
         Layers for the GMM network wrapped in a keras.Sequential class.
-    n_gmm
+    n_gmm : int
         Number of components in GMM.
-    latent_dim
+    latent_dim : int
         Dimensionality of the latent space.
-    recon_features
+    recon_features : Callable, default eucl_cosim_features
         Function to extract features from the reconstructed instance by the decoder.
-    beta
+    beta : float, default 1.0
         Beta parameter for KL-divergence loss term.
     """
 
@@ -305,8 +308,8 @@ class VAEGMM(keras.Model):
         self.recon_features = recon_features
         self.beta = beta
 
-    def call(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
-        enc_mean, enc_log_var, enc = cast(Tuple[tf.Tensor, tf.Tensor, tf.Tensor], self.encoder(x))
+    def call(self, x: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        enc_mean, enc_log_var, enc = cast(tuple[tf.Tensor, tf.Tensor, tf.Tensor], self.encoder(x))
         x_recon = cast(tf.Tensor, self.decoder(enc))
         recon_features = self.recon_features(x, x_recon)
         z = cast(tf.Tensor, tf.concat([enc, recon_features], -1))

dataeval/_internal/models/tensorflow/gmm.py

@@ -6,7 +6,9 @@ Original code Copyright (c) 2023 Seldon Technologies Ltd
 Licensed under Apache Software License (Apache 2.0)
 """
 
-from typing import NamedTuple, Tuple
+from __future__ import annotations
+
+from typing import NamedTuple
 
 import numpy as np
 import tensorflow as tf
@@ -75,7 +77,7 @@ def gmm_energy(
     z: tf.Tensor,
     params: GaussianMixtureModelParams,
     return_mean: bool = True,
-) -> Tuple[tf.Tensor, tf.Tensor]:
+) -> tuple[tf.Tensor, tf.Tensor]:
     """
     Compute sample energy from Gaussian Mixture Model.
 

dataeval/_internal/models/tensorflow/losses.py

@@ -6,7 +6,9 @@ Original code Copyright (c) 2023 Seldon Technologies Ltd
 Licensed under Apache Software License (Apache 2.0)
 """
 
-from typing import Literal, Optional, Union, cast
+from __future__ import annotations
+
+from typing import Literal, cast
 
 import tensorflow as tf
 from keras.layers import Flatten
@@ -20,22 +22,24 @@ from dataeval._internal.models.tensorflow.gmm import gmm_energy, gmm_params
 
 class Elbo:
     """
-    Compute ELBO loss. The covariance matrix can be specified by passing the full covariance matrix, the matrix
+    Compute ELBO loss.
+
+    The covariance matrix can be specified by passing the full covariance matrix, the matrix
     diagonal, or a scale identity multiplier. Only one of these should be specified. If none are specified, the
     identity matrix is used.
 
     Parameters
     ----------
-    cov_type
+    cov_type : Union[Literal["cov_full", "cov_diag"], float], default 1.0
         Full covariance matrix, diagonal variance matrix, or scale identity multiplier.
-    x
+    x : ArrayLike, optional - default None
         Dataset used to calculate the covariance matrix.  Required for full and diagonal covariance matrix types.
     """
 
     def __init__(
         self,
-        cov_type: Union[Literal["cov_full", "cov_diag"], float] = 1.0,
-        x: Optional[Union[tf.Tensor, NDArray]] = None,
+        cov_type: Literal["cov_full", "cov_diag"] | float = 1.0,
+        x: tf.Tensor | NDArray | None = None,
     ):
         if isinstance(cov_type, float):
             self.cov = ("sim", cov_type)
@@ -67,13 +71,13 @@ class LossGMM:
 
     Parameters
     ----------
-    w_recon
+    w_recon : float, default 1e-7
         Weight on elbo loss term.
-    w_energy
+    w_energy : float, default 0.1
         Weight on sample energy loss term.
-    w_cov_diag
+    w_cov_diag : float, default 0.005
         Weight on covariance regularizing loss term.
-    elbo
+    elbo : Elbo, optional - default None
         ELBO loss function used to calculate w_recon.
     """
 
@@ -82,7 +86,7 @@ class LossGMM:
         w_recon: float = 1e-7,
         w_energy: float = 0.1,
         w_cov_diag: float = 0.005,
-        elbo: Optional[Elbo] = None,
+        elbo: Elbo | None = None,
     ):
         self.w_recon = w_recon
         self.w_energy = w_energy

dataeval/_internal/models/tensorflow/pixelcnn.py

@@ -8,9 +8,10 @@ Original code Copyright (c) 2023 Seldon Technologies Ltd
 Licensed under Apache Software License (Apache 2.0)
 """
 
+from __future__ import annotations
+
 import functools
 import warnings
-from typing import Optional
 
 import keras
 import numpy as np
@@ -238,47 +239,47 @@ class PixelCNN(distribution.Distribution):
 
     Parameters
     ----------
-    image_shape
+    image_shape : tuple
         3D `TensorShape` or tuple for the `[height, width, channels]` dimensions of the image.
-    conditional_shape
+    conditional_shape : tuple, optional - default None
         `TensorShape` or tuple for the shape of the conditional input, or `None` if there is no conditional input.
-    num_resnet
+    num_resnet : int, default 5
         The number of layers (shown in Figure 2 of [2]) within each highest-level block of Figure 2 of [1].
-    num_hierarchies
+    num_hierarchies : int, default 3
         The number of highest-level blocks (separated by expansions/contractions of dimensions in Figure 2 of [1].)
-    num_filters
+    num_filters : int, default 160
         The number of convolutional filters.
-    num_logistic_mix
+    num_logistic_mix : int, default 10
         Number of components in the logistic mixture distribution.
-    receptive_field_dims
+    receptive_field_dims tuple, default (3, 3)
         Height and width in pixels of the receptive field of the convolutional layers above and to the left
         of a given pixel. The width (second element of the tuple) should be odd. Figure 1 (middle) of [2]
         shows a receptive field of (3, 5) (the row containing the current pixel is included in the height).
         The default of (3, 3) was used to produce the results in [1].
-    dropout_p
+    dropout_p : float, default 0.0
         The dropout probability. Should be between 0 and 1.
-    resnet_activation
+    resnet_activation : str, default "concat_elu"
         The type of activation to use in the resnet blocks. May be 'concat_elu', 'elu', or 'relu'.
-    l2_weight
+    l2_weight : float, default 0.0
         The L2 regularization weight.
-    use_weight_norm
+    use_weight_norm : bool, default True
         If `True` then use weight normalization (works only in Eager mode).
-    use_data_init
+    use_data_init : bool, default True
         If `True` then use data-dependent initialization (has no effect if `use_weight_norm` is `False`).
-    high
+    high : int, default 255
         The maximum value of the input data (255 for an 8-bit image).
-    low
+    low : int, default 0
         The minimum value of the input data.
-    dtype
+    dtype : tensorflow dtype, default tf.float32
         Data type of the `Distribution`.
-    name
+    name : str, default "PixelCNN"
         The name of the `Distribution`.
     """
 
     def __init__(
         self,
         image_shape: tuple,
-        conditional_shape: Optional[tuple] = None,
+        conditional_shape: tuple | None = None,
         num_resnet: int = 5,
         num_hierarchies: int = 3,
         num_filters: int = 160,

dataeval/_internal/models/tensorflow/trainer.py

@@ -6,7 +6,9 @@ Original code Copyright (c) 2023 Seldon Technologies Ltd
 Licensed under Apache Software License (Apache 2.0)
 """
 
-from typing import Callable, Iterable, Optional, Tuple, cast
+from __future__ import annotations
+
+from typing import Callable, Iterable, cast
 
 import keras
 import numpy as np
@@ -17,10 +19,10 @@ from numpy.typing import NDArray
 def trainer(
     model: keras.Model,
     x_train: NDArray,
-    y_train: Optional[NDArray] = None,
-    loss_fn: Optional[Callable[..., tf.Tensor]] = None,
+    y_train: NDArray | None = None,
+    loss_fn: Callable[..., tf.Tensor] | None = None,
     optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
-    preprocess_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None,
+    preprocess_fn: Callable[[tf.Tensor], tf.Tensor] | None = None,
     epochs: int = 20,
     reg_loss_fn: Callable[[keras.Model], tf.Tensor] = (lambda _: cast(tf.Tensor, tf.Variable(0, dtype=tf.float32))),
     batch_size: int = 64,
@@ -70,14 +72,14 @@ def trainer(
             dataset.on_epoch_end()  # type: ignore py39
         loss_val_ma = 0.0
         for step, data in enumerate(dataset):
-            x, y = cast(Tuple[tf.Tensor, Optional[tf.Tensor]], data if isinstance(data, tuple) else (data, None))
+            x, y = data if isinstance(data, tuple) else (data, None)
             if isinstance(preprocess_fn, Callable):
                 x = preprocess_fn(x)
             with tf.GradientTape() as tape:
                 y_hat = model(x)
                 y = x if y is None else y
                 if isinstance(loss_fn, Callable):
-                    args = [y] + list(y_hat) if isinstance(y_hat, Tuple) else [y, y_hat]
+                    args = [y] + list(y_hat) if isinstance(y_hat, tuple) else [y, y_hat]
                     loss = loss_fn(*args)
                 else:
                     loss = cast(tf.Tensor, tf.constant(0.0, dtype=tf.float32))