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

dataeval/{_internal/models/tensorflow → utils/tensorflow/_internal}/autoencoder.py

@@ -6,14 +6,14 @@ Original code Copyright (c) 2023 Seldon Technologies Ltd
 Licensed under Apache Software License (Apache 2.0)
 """
 
-# pyright: reportIncompatibleMethodOverride=false
-
 from __future__ import annotations
 
-from typing import Callable, cast
+from typing import cast
 
 import tensorflow as tf
 import tf_keras as keras
+from tensorflow.python.module.module import Module  # noqa
+from tf_keras import Sequential
 from tf_keras.layers import (
     Dense,
     Flatten,
@@ -90,7 +90,7 @@ class Sampling(Layer):
         Parameters
         ----------
         inputs
-            Tuple with mean and log variance.
+            Tuple with mean and log :term:`variance<Variance>`.
 
         Returns
         -------
@@ -103,7 +103,7 @@ class Sampling(Layer):
 
 
 class EncoderAE(Layer):
-    def __init__(self, encoder_net: keras.Model) -> None:
+    def __init__(self, encoder_net: Sequential) -> None:
         """
         Encoder of AE.
 
@@ -115,14 +115,14 @@ class EncoderAE(Layer):
             Name of encoder.
         """
         super().__init__(name="encoder_ae")
-        self.encoder_net = encoder_net
+        self.encoder_net: Sequential = encoder_net
 
     def call(self, x: tf.Tensor) -> tf.Tensor:
         return cast(tf.Tensor, self.encoder_net(x))
 
 
 class EncoderVAE(Layer):
-    def __init__(self, encoder_net: keras.Model, latent_dim: int) -> None:
+    def __init__(self, encoder_net: Sequential, latent_dim: int) -> None:
         """
         Encoder of VAE.
 
@@ -131,28 +131,28 @@ class EncoderVAE(Layer):
         encoder_net
             Layers for the encoder wrapped in a keras.Sequential class.
         latent_dim
-            Dimensionality of the latent space.
+            Dimensionality of the :term:`latent space<Latent Space>`.
         name
             Name of encoder.
         """
         super().__init__(name="encoder_vae")
-        self.encoder_net = encoder_net
-        self.fc_mean = Dense(latent_dim, activation=None)
-        self.fc_log_var = Dense(latent_dim, activation=None)
-        self.sampling = Sampling()
+        self.encoder_net: Sequential = encoder_net
+        self._fc_mean = Dense(latent_dim, activation=None)
+        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]:
         x = cast(tf.Tensor, self.encoder_net(x))
         if len(x.shape) > 2:
             x = cast(tf.Tensor, Flatten()(x))
-        z_mean = cast(tf.Tensor, self.fc_mean(x))
-        z_log_var = cast(tf.Tensor, self.fc_log_var(x))
-        z = cast(tf.Tensor, self.sampling((z_mean, z_log_var)))
+        z_mean = cast(tf.Tensor, self._fc_mean(x))
+        z_log_var = cast(tf.Tensor, self._fc_log_var(x))
+        z = cast(tf.Tensor, self._sampling((z_mean, z_log_var)))
         return z_mean, z_log_var, z
 
 
 class Decoder(Layer):
-    def __init__(self, decoder_net: keras.Model) -> None:
+    def __init__(self, decoder_net: Sequential) -> None:
         """
         Decoder of AE and VAE.
 
@@ -164,10 +164,10 @@ class Decoder(Layer):
             Name of decoder.
         """
         super().__init__(name="decoder")
-        self.decoder_net = decoder_net
+        self.decoder_net: Sequential = decoder_net
 
-    def call(self, x: tf.Tensor) -> tf.Tensor:
-        return cast(tf.Tensor, self.decoder_net(x))
+    def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        return cast(tf.Tensor, self.decoder_net(inputs))
 
 
 class AE(keras.Model):
@@ -176,19 +176,19 @@ class AE(keras.Model):
 
     Parameters
     ----------
-    encoder_net : keras.Model
+    encoder_net : Sequential
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net : keras.Model
+    decoder_net : Sequential
         Layers for the decoder wrapped in a keras.Sequential class.
     """
 
-    def __init__(self, encoder_net: keras.Model, decoder_net: keras.Model) -> None:
+    def __init__(self, encoder_net: Sequential, decoder_net: Sequential) -> None:
         super().__init__(name="ae")
-        self.encoder = EncoderAE(encoder_net)
-        self.decoder = Decoder(decoder_net)
+        self.encoder: Layer = EncoderAE(encoder_net)
+        self.decoder: Layer = Decoder(decoder_net)
 
-    def call(self, x: tf.Tensor) -> tf.Tensor:
-        z = cast(tf.Tensor, self.encoder(x))
+    def call(self, inputs: tf.Tensor, training: bool | None = None, mask: tf.Tensor | None = None) -> tf.Tensor:
+        z = cast(tf.Tensor, self.encoder(inputs))
         x_recon = cast(tf.Tensor, self.decoder(z))
         return x_recon
 
@@ -199,25 +199,25 @@ class VAE(keras.Model):
 
     Parameters
     ----------
-    encoder_net : keras.Model
+    encoder_net : Sequential
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net : keras.Model
+    decoder_net : Sequential
         Layers for the decoder wrapped in a keras.Sequential class.
     latent_dim : int
-        Dimensionality of the latent space.
+        Dimensionality of the :term:`latent space<Latent Space>`.
     beta : float, default 1.0
         Beta parameter for KL-divergence loss term.
     """
 
-    def __init__(self, encoder_net: keras.Model, decoder_net: keras.Model, latent_dim: int, beta: float = 1.0) -> None:
+    def __init__(self, encoder_net: Sequential, decoder_net: Sequential, latent_dim: int, beta: float = 1.0) -> None:
         super().__init__(name="vae_model")
-        self.encoder = EncoderVAE(encoder_net, latent_dim)
-        self.decoder = Decoder(decoder_net)
-        self.beta = beta
-        self.latent_dim = latent_dim
+        self.encoder: Layer = EncoderVAE(encoder_net, latent_dim)
+        self.decoder: Layer = Decoder(decoder_net)
+        self.beta: float = beta
+        self.latent_dim: int = 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))
+    def call(self, inputs: tf.Tensor, training: bool | None = None, mask: tf.Tensor | None = None) -> tf.Tensor:
+        z_mean, z_log_var, z = cast(tuple[tf.Tensor, tf.Tensor, tf.Tensor], self.encoder(inputs))
         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)
@@ -231,37 +231,35 @@ class AEGMM(keras.Model):
 
     Parameters
     ----------
-    encoder_net : keras.Model
+    encoder_net : Sequential
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net : keras.Model
+    decoder_net : Sequential
         Layers for the decoder wrapped in a keras.Sequential class.
-    gmm_density_net : keras.Model
+    gmm_density_net : Sequential
         Layers for the GMM network wrapped in a keras.Sequential class.
     n_gmm : int
         Number of components in GMM.
-    recon_features : Callable, default eucl_cosim_features
-        Function to extract features from the reconstructed instance by the decoder.
     """
 
     def __init__(
         self,
-        encoder_net: keras.Model,
-        decoder_net: keras.Model,
-        gmm_density_net: keras.Model,
+        encoder_net: Sequential,
+        decoder_net: Sequential,
+        gmm_density_net: Sequential,
         n_gmm: int,
-        recon_features: Callable = eucl_cosim_features,
     ) -> None:
         super().__init__("aegmm")
         self.encoder = encoder_net
         self.decoder = decoder_net
         self.gmm_density = gmm_density_net
         self.n_gmm = n_gmm
-        self.recon_features = recon_features
 
-    def call(self, x: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
-        enc = self.encoder(x)
+    def call(
+        self, inputs: tf.Tensor, training: bool | None = None, mask: tf.Tensor | None = None
+    ) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        enc = self.encoder(inputs)
         x_recon = cast(tf.Tensor, self.decoder(enc))
-        recon_features = self.recon_features(x, x_recon)
+        recon_features = eucl_cosim_features(inputs, x_recon)
         z = cast(tf.Tensor, tf.concat([enc, recon_features], -1))
         gamma = cast(tf.Tensor, self.gmm_density(z))
         return x_recon, z, gamma
@@ -273,45 +271,43 @@ class VAEGMM(keras.Model):
 
     Parameters
     ----------
-    encoder_net : keras.Model
+    encoder_net : Sequential
         Layers for the encoder wrapped in a keras.Sequential class.
-    decoder_net : keras.Model
+    decoder_net : Sequential
         Layers for the decoder wrapped in a keras.Sequential class.
-    gmm_density_net : keras.Model
+    gmm_density_net : Sequential
         Layers for the GMM network wrapped in a keras.Sequential class.
     n_gmm : int
         Number of components in GMM.
     latent_dim : int
-        Dimensionality of the latent space.
-    recon_features : Callable, default eucl_cosim_features
-        Function to extract features from the reconstructed instance by the decoder.
+        Dimensionality of the :term:`latent space<Latent Space>`.
     beta : float, default 1.0
         Beta parameter for KL-divergence loss term.
     """
 
     def __init__(
         self,
-        encoder_net: keras.Model,
-        decoder_net: keras.Model,
-        gmm_density_net: keras.Model,
+        encoder_net: Sequential,
+        decoder_net: Sequential,
+        gmm_density_net: Sequential,
         n_gmm: int,
         latent_dim: int,
-        recon_features: Callable = eucl_cosim_features,
         beta: float = 1.0,
     ) -> None:
         super().__init__(name="vaegmm")
-        self.encoder = EncoderVAE(encoder_net, latent_dim)
-        self.decoder = decoder_net
-        self.gmm_density = gmm_density_net
-        self.n_gmm = n_gmm
-        self.latent_dim = latent_dim
-        self.recon_features = recon_features
+        self.encoder: Sequential = EncoderVAE(encoder_net, latent_dim)
+        self.decoder: Sequential = decoder_net
+        self.gmm_density: Sequential = gmm_density_net
+        self.n_gmm: int = n_gmm
+        self.latent_dim: int = latent_dim
         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, inputs: tf.Tensor, training: bool | None = None, mask: tf.Tensor | None = None
+    ) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        enc_mean, enc_log_var, enc = cast(tuple[tf.Tensor, tf.Tensor, tf.Tensor], self.encoder(inputs))
         x_recon = cast(tf.Tensor, self.decoder(enc))
-        recon_features = self.recon_features(x, x_recon)
+        recon_features = eucl_cosim_features(inputs, x_recon)
         z = cast(tf.Tensor, tf.concat([enc, recon_features], -1))
         gamma = cast(tf.Tensor, self.gmm_density(z))
         # add KL divergence loss term

dataeval/{_internal/models/tensorflow/losses.py → utils/tensorflow/_internal/loss.py}

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 from typing import Literal, cast
 
+import numpy as np
 import tensorflow as tf
 from numpy.typing import NDArray
 from tensorflow_probability.python.distributions.mvn_diag import MultivariateNormalDiag
@@ -17,7 +18,7 @@ from tensorflow_probability.python.distributions.mvn_tril import MultivariateNor
 from tensorflow_probability.python.stats import covariance
 from tf_keras.layers import Flatten
 
-from dataeval._internal.models.tensorflow.gmm import gmm_energy, gmm_params
+from dataeval.utils.tensorflow._internal.gmm import gmm_energy, gmm_params
 
 
 class Elbo:
@@ -31,7 +32,7 @@ class Elbo:
     Parameters
     ----------
     cov_type : Union[Literal["cov_full", "cov_diag"], float], default 1.0
-        Full covariance matrix, diagonal variance matrix, or scale identity multiplier.
+        Full covariance matrix, diagonal :term:`variance<Variance>` matrix, or scale identity multiplier.
     x : ArrayLike, optional - default None
         Dataset used to calculate the covariance matrix.  Required for full and diagonal covariance matrix types.
     """
@@ -39,26 +40,26 @@ class Elbo:
     def __init__(
         self,
         cov_type: Literal["cov_full", "cov_diag"] | float = 1.0,
-        x: tf.Tensor | NDArray | None = None,
+        x: tf.Tensor | NDArray[np.float32] | None = None,
     ):
         if isinstance(cov_type, float):
-            self.cov = ("sim", cov_type)
+            self._cov = ("sim", cov_type)
         elif cov_type in ["cov_full", "cov_diag"]:
-            x_np: NDArray = x.numpy() if tf.is_tensor(x) else x  # type: ignore
+            x_np: NDArray[np.float32] = x.numpy().astype(np.float32) if tf.is_tensor(x) else x  # type: ignore
             cov = covariance(x_np.reshape(x_np.shape[0], -1))  # type: ignore py38
             if cov_type == "cov_diag":  # infer standard deviation from covariance matrix
                 cov = tf.math.sqrt(tf.linalg.diag_part(cov))
-            self.cov = (cov_type, cov)
+            self._cov = (cov_type, cov)
         else:
             raise ValueError("Only cov_full, cov_diag or sim value should be specified.")
 
     def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         y_pred_flat = cast(tf.Tensor, Flatten()(y_pred))
 
-        if self.cov[0] == "cov_full":
-            y_mn = MultivariateNormalTriL(y_pred_flat, scale_tril=tf.linalg.cholesky(self.cov[1]))
+        if self._cov[0] == "cov_full":
+            y_mn = MultivariateNormalTriL(y_pred_flat, scale_tril=tf.linalg.cholesky(self._cov[1]))
         else:  # cov_diag and sim
-            cov_diag = self.cov[1] if self.cov[0] == "cov_diag" else self.cov[1] * tf.ones(y_pred_flat.shape[-1])
+            cov_diag = self._cov[1] if self._cov[0] == "cov_diag" else self._cov[1] * tf.ones(y_pred_flat.shape[-1])
             y_mn = MultivariateNormalDiag(y_pred_flat, scale_diag=cov_diag)
 
         loss = -tf.reduce_mean(y_mn.log_prob(Flatten()(y_true)))

dataeval/{_internal/models/tensorflow → utils/tensorflow/_internal}/pixelcnn.py

@@ -34,13 +34,9 @@ from tensorflow_probability.python.internal import (
     tensorshape_util,
 )
 
-__all__ = [
-    "Shift",
-]
-
 
 class WeightNorm(keras.layers.Wrapper):
-    def __init__(self, layer, data_init: bool = True, **kwargs):
+    def __init__(self, layer, data_init: bool = True, **kwargs) -> None:
         """Layer wrapper to decouple magnitude and direction of the layer's weights.
 
         This wrapper reparameterizes a layer by decoupling the weight's
@@ -187,7 +183,7 @@ class WeightNorm(keras.layers.Wrapper):
 
 
 class Shift(bijector.Bijector):
-    def __init__(self, shift, validate_args=False, name="shift"):
+    def __init__(self, shift, validate_args=False, name="shift") -> None:
         """Instantiates the `Shift` bijector which computes `Y = g(X; shift) = X + shift`
         where `shift` is a numeric `Tensor`.
 
@@ -276,13 +272,13 @@ class PixelCNN(distribution.Distribution):
 
     def __init__(
         self,
-        image_shape: tuple,
-        conditional_shape: tuple | None = None,
+        image_shape: tuple[int, int, int],
+        conditional_shape: tuple[int, ...] | None = None,
         num_resnet: int = 5,
         num_hierarchies: int = 3,
         num_filters: int = 160,
         num_logistic_mix: int = 10,
-        receptive_field_dims: tuple = (3, 3),
+        receptive_field_dims: tuple[int, int] = (3, 3),
         dropout_p: float = 0.5,
         resnet_activation: str = "concat_elu",
         l2_weight: float = 0.0,
@@ -290,7 +286,7 @@ class PixelCNN(distribution.Distribution):
         use_data_init: bool = True,
         high: int = 255,
         low: int = 0,
-        dtype=tf.float32,
+        dtype: tf.DType = tf.float32,
     ) -> None:
         parameters = dict(locals())
         with tf.name_scope("PixelCNN") as name:
@@ -315,7 +311,7 @@ class PixelCNN(distribution.Distribution):
             self._high = tf.cast(high, self.dtype)
             self._low = tf.cast(low, self.dtype)
             self._num_logistic_mix = num_logistic_mix
-            self.network = _PixelCNNNetwork(
+            self._network = PixelCNNNetwork(
                 dropout_p=dropout_p,
                 num_resnet=num_resnet,
                 num_hierarchies=num_hierarchies,
@@ -338,7 +334,7 @@ class PixelCNN(distribution.Distribution):
 
             self.image_shape = image_shape
             self.conditional_shape = conditional_shape
-            self.network.build(input_shape)
+            self._network.build(input_shape)
 
     def _make_mixture_dist(self, component_logits, locs, scales, return_per_feature: bool = False):
         """Builds a mixture of quantized logistic distributions.
@@ -400,7 +396,7 @@ class PixelCNN(distribution.Distribution):
         Parameters
         ----------
         value
-            `Tensor` or Numpy array of image data. May have leading batch
+            `Tensor` or :term:`NumPy` array of image data. May have leading batch
             dimension(s), which must broadcast to the leading batch dimensions of
             `conditional_input`.
         conditional_input
@@ -455,7 +451,7 @@ class PixelCNN(distribution.Distribution):
         transformed_value = (2.0 * (value - self._low) / (self._high - self._low)) - 1.0
         inputs = transformed_value if conditional_input is None else [transformed_value, conditional_input]
 
-        params = self.network(inputs, training=training)
+        params = self._network(inputs, training=training)
 
         num_channels = self.event_shape[-1]
         if num_channels == 1:
@@ -554,7 +550,7 @@ class PixelCNN(distribution.Distribution):
             seed=seed,
         )
         inputs = samples_0 if conditional_input is None else [samples_0, h]
-        params_0 = self.network(inputs, training=training)
+        params_0 = self._network(inputs, training=training)
         samples_0 = self._sample_channels(*params_0, seed=seed)
 
         image_height, image_width, _ = tensorshape_util.as_list(self.event_shape)
@@ -579,7 +575,7 @@ class PixelCNN(distribution.Distribution):
                 width, num_channels]`.
             """
             inputs = samples if conditional_input is None else [samples, h]
-            params = self.network(inputs, training=training)
+            params = self._network(inputs, training=training)
             samples_new = self._sample_channels(*params, seed=seed)
 
             # Update the current pixel
@@ -673,7 +669,7 @@ class PixelCNN(distribution.Distribution):
         return tf.TensorShape(self.image_shape)
 
 
-class _PixelCNNNetwork(keras.layers.Layer):
+class PixelCNNNetwork(keras.layers.Layer):
     """Keras `Layer` to parameterize a Pixel CNN++ distribution.
     This is a Keras implementation of the Pixel CNN++ network, as described in
     Salimans et al. (2017)[1] and van den Oord et al. (2016)[2].
@@ -699,14 +695,14 @@ class _PixelCNNNetwork(keras.layers.Layer):
         num_hierarchies: int = 3,
         num_filters: int = 160,
         num_logistic_mix: int = 10,
-        receptive_field_dims: tuple = (3, 3),
+        receptive_field_dims: tuple[int, int] = (3, 3),
         resnet_activation: str = "concat_elu",
         l2_weight: float = 0.0,
         use_weight_norm: bool = True,
         use_data_init: bool = True,
-        dtype=tf.float32,
+        dtype: tf.DType = tf.float32,
     ) -> None:
-        """Initialize the neural network for the Pixel CNN++ distribution.
+        """Initialize the :term:`neural network<Neural Network>` for the Pixel CNN++ distribution.
 
         Parameters
         ----------
@@ -765,7 +761,7 @@ class _PixelCNNNetwork(keras.layers.Layer):
         else:
             self._layer_wrapper = lambda layer: layer
 
-    def build(self, input_shape):
+    def build(self, input_shape: tuple[int, ...]) -> None:
         dtype = self.dtype
         if len(input_shape) == 2:
             batch_image_shape, batch_conditional_shape = input_shape
@@ -1040,7 +1036,7 @@ class _PixelCNNNetwork(keras.layers.Layer):
         self._network = keras.Model(inputs=inputs, outputs=outputs)
         super().build(input_shape)
 
-    def call(self, inputs, training=None):
+    def call(self, inputs: tf.Tensor, training: bool | None = None, mask: tf.Tensor | None = None) -> tf.Tensor:
         """Call the Pixel CNN network model.
 
         Parameters

dataeval/{_internal/models/tensorflow → utils/tensorflow/_internal}/trainer.py

@@ -60,7 +60,9 @@ def trainer(
     loss_fn = loss_fn() if isinstance(loss_fn, type) else loss_fn
     optimizer = optimizer() if isinstance(optimizer, type) else optimizer
 
-    train_data = x_train if y_train is None else (x_train, y_train)
+    train_data = (
+        x_train.astype(np.float32) if y_train is None else (x_train.astype(np.float32), y_train.astype(np.float32))
+    )
     dataset = tf.data.Dataset.from_tensor_slices(train_data)
     dataset = dataset.shuffle(buffer_size=buffer_size).batch(batch_size)
     n_minibatch = len(dataset)

dataeval/{_internal/models/tensorflow → utils/tensorflow/_internal}/utils.py

@@ -9,7 +9,7 @@ Licensed under Apache Software License (Apache 2.0)
 from __future__ import annotations
 
 import math
-from typing import Callable, Union, cast
+from typing import Any, Callable, Literal, Union, cast
 
 import numpy as np
 import tensorflow as tf
@@ -26,8 +26,8 @@ from tf_keras.layers import (
     Reshape,
 )
 
-from dataeval._internal.models.tensorflow.autoencoder import AE, AEGMM, VAE, VAEGMM
-from dataeval._internal.models.tensorflow.pixelcnn import PixelCNN
+from dataeval.utils.tensorflow._internal.autoencoder import AE, AEGMM, VAE, VAEGMM
+from dataeval.utils.tensorflow._internal.pixelcnn import PixelCNN
 
 
 def predict_batch(
@@ -55,7 +55,7 @@ def predict_batch(
 
     Returns
     -------
-    Numpy array, tensorflow tensor or tuples of those with model outputs.
+    :term:`NumPy` array, tensorflow tensor or tuples of those with model outputs.
     """
     n = len(x)
     n_minibatch = int(np.ceil(n / batch_size))
@@ -95,7 +95,7 @@ def predict_batch(
     return out
 
 
-def _get_default_encoder_net(input_shape: tuple[int, int, int], encoding_dim: int):
+def get_default_encoder_net(input_shape: tuple[int, int, int], encoding_dim: int):
     return Sequential(
         [
             InputLayer(input_shape=input_shape),
@@ -108,7 +108,7 @@ def _get_default_encoder_net(input_shape: tuple[int, int, int], encoding_dim: in
     )
 
 
-def _get_default_decoder_net(input_shape: tuple[int, int, int], encoding_dim: int):
+def get_default_decoder_net(input_shape: tuple[int, int, int], encoding_dim: int):
     return Sequential(
         [
             InputLayer(input_shape=(encoding_dim,)),
@@ -124,18 +124,18 @@ def _get_default_decoder_net(input_shape: tuple[int, int, int], encoding_dim: in
 
 
 def create_model(
-    model_type: AE | AEGMM | PixelCNN | VAE | VAEGMM,
+    model_type: Literal["AE", "AEGMM", "PixelCNN", "VAE", "VAEGMM"],
     input_shape: tuple[int, int, int],
     encoding_dim: int | None = None,
     n_gmm: int | None = None,
     gmm_latent_dim: int | None = None,
-):
+) -> Any:
     """
     Create a default model for the specified model type.
 
     Parameters
     ----------
-    model_type : Union[AE, AEGMM, PixelCNN, VAE, VAEGMM]
+    model_type : Literal["AE", "AEGMM", "PixelCNN", "VAE", "VAEGMM"]
         The model type to create.
     input_shape : Tuple[int, int, int]
         The input shape of the data used.
@@ -148,20 +148,20 @@ def create_model(
     """
     input_dim = math.prod(input_shape)
     encoding_dim = int(math.pow(2, int(input_dim.bit_length() * 0.8)) if encoding_dim is None else encoding_dim)
-    if model_type == AE:
+    if model_type == "AE":
         return AE(
-            _get_default_encoder_net(input_shape, encoding_dim),
-            _get_default_decoder_net(input_shape, encoding_dim),
+            get_default_encoder_net(input_shape, encoding_dim),
+            get_default_decoder_net(input_shape, encoding_dim),
         )
 
-    if model_type == VAE:
+    if model_type == "VAE":
         return VAE(
-            _get_default_encoder_net(input_shape, encoding_dim),
-            _get_default_decoder_net(input_shape, encoding_dim),
+            get_default_encoder_net(input_shape, encoding_dim),
+            get_default_decoder_net(input_shape, encoding_dim),
             encoding_dim,
         )
 
-    if model_type == AEGMM:
+    if model_type == "AEGMM":
         n_gmm = 2 if n_gmm is None else n_gmm
         gmm_latent_dim = 1 if gmm_latent_dim is None else gmm_latent_dim
         # The outlier detector is an encoder/decoder architecture
@@ -201,7 +201,7 @@ def create_model(
             n_gmm=n_gmm,
         )
 
-    if model_type == VAEGMM:
+    if model_type == "VAEGMM":
         n_gmm = 2 if n_gmm is None else n_gmm
         gmm_latent_dim = 2 if gmm_latent_dim is None else gmm_latent_dim
         # The outlier detector is an encoder/decoder architecture
@@ -242,7 +242,7 @@ def create_model(
             latent_dim=gmm_latent_dim,
         )
 
-    if model_type == PixelCNN:
+    if model_type == "PixelCNN":
         return PixelCNN(
             image_shape=input_shape,
             num_resnet=5,

dataeval/utils/tensorflow/loss/__init__.py

@@ -1,7 +1,11 @@
 from dataeval import _IS_TENSORFLOW_AVAILABLE
-from dataeval._internal.models.tensorflow.losses import Elbo, LossGMM
 
 __all__ = []
 
+
 if _IS_TENSORFLOW_AVAILABLE:
-    __all__ += ["Elbo", "LossGMM"]
+    from dataeval.utils.tensorflow._internal.loss import Elbo, LossGMM
+
+    __all__ = ["Elbo", "LossGMM"]
+
+del _IS_TENSORFLOW_AVAILABLE

dataeval/utils/torch/__init__.py

@@ -6,16 +6,20 @@ to create a seamless integration between custom models and DataEval's metrics.
 """
 
 from dataeval import _IS_TORCH_AVAILABLE, _IS_TORCHVISION_AVAILABLE
-from dataeval._internal.utils import read_dataset
 
 __all__ = []
 
 if _IS_TORCH_AVAILABLE:
-    from . import models, trainer
+    from dataeval.utils.torch import models, trainer
+    from dataeval.utils.torch.utils import read_dataset
 
     __all__ += ["read_dataset", "models", "trainer"]
 
 if _IS_TORCHVISION_AVAILABLE:
-    from . import datasets
+    from dataeval.utils.torch import datasets
 
     __all__ += ["datasets"]
+
+
+del _IS_TORCH_AVAILABLE
+del _IS_TORCHVISION_AVAILABLE