diff --git a/nalp/models/dcgan.py b/nalp/models/dcgan.py index 730310d..ca53cba 100644 --- a/nalp/models/dcgan.py +++ b/nalp/models/dcgan.py @@ -1,6 +1,8 @@ """Deep Convolutional Generative Adversarial Network. """ +from typing import Optional, Tuple + from nalp.core import Adversarial from nalp.models.discriminators import ConvDiscriminator from nalp.models.generators import ConvGenerator @@ -21,20 +23,20 @@ class DCGAN(Adversarial): def __init__( self, - input_shape=(28, 28, 1), - noise_dim=100, - n_samplings=3, - alpha=0.3, - dropout_rate=0.3, - ): + input_shape: Optional[Tuple[int, int, int]] = (28, 28, 1), + noise_dim: Optional[int] = 100, + n_samplings: Optional[int] = 3, + alpha: Optional[float] = 0.3, + dropout_rate: Optional[float] = 0.3, + ) -> None: """Initialization method. Args: - input_shape (tuple): An input shape for the Generator. - noise_dim (int): Amount of noise dimensions for the Generator. - n_samplings (int): Number of down/up samplings to perform. - alpha (float): LeakyReLU activation threshold. - dropout_rate (float): Dropout activation rate. + input_shape: An input shape for the Generator. + noise_dim: Amount of noise dimensions for the Generator. + n_samplings: Number of down/up samplings to perform. + alpha: LeakyReLU activation threshold. + dropout_rate: Dropout activation rate. """ diff --git a/nalp/models/gan.py b/nalp/models/gan.py index 6b0085c..18db683 100644 --- a/nalp/models/gan.py +++ b/nalp/models/gan.py @@ -1,6 +1,8 @@ """Generative Adversarial Network. """ +from typing import Optional, Tuple + from nalp.core import Adversarial from nalp.models.discriminators import LinearDiscriminator from nalp.models.generators import LinearGenerator @@ -18,14 +20,20 @@ class GAN(Adversarial): """ - def __init__(self, input_shape=(784,), noise_dim=100, n_samplings=3, alpha=0.01): + def __init__( + self, + input_shape: Optional[Tuple[int, ...]] = (784,), + noise_dim: Optional[int] = 100, + n_samplings: Optional[int] = 3, + alpha: Optional[float] = 0.01, + ) -> None: """Initialization method. Args: - input_shape (tuple): An input shape for the Generator. - noise_dim (int): Amount of noise dimensions for the Generator. - n_samplings (int): Number of down/up samplings to perform. - alpha (float): LeakyReLU activation threshold. + input_shape: An input shape for the Generator. + noise_dim: Amount of noise dimensions for the Generator. + n_samplings: Number of down/up samplings to perform. + alpha: LeakyReLU activation threshold. """ diff --git a/nalp/models/gsgan.py b/nalp/models/gsgan.py index 55f52da..e9f7287 100644 --- a/nalp/models/gsgan.py +++ b/nalp/models/gsgan.py @@ -1,10 +1,14 @@ """Gumbel-Softmax Generative Adversarial Network. """ +from typing import Optional, Tuple + import tensorflow as tf from tensorflow.keras.utils import Progbar from nalp.core import Adversarial +from nalp.core.dataset import Dataset +from nalp.encoders.integer import IntegerEncoder from nalp.models.discriminators import LSTMDiscriminator from nalp.models.generators import GumbelLSTMGenerator from nalp.utils import logging @@ -24,16 +28,21 @@ class GSGAN(Adversarial): """ def __init__( - self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64, tau=5 - ): + self, + encoder: Optional[IntegerEncoder] = None, + vocab_size: Optional[int] = 1, + embedding_size: Optional[int] = 32, + hidden_size: Optional[int] = 64, + tau: Optional[float] = 5, + ) -> None: """Initialization method. Args: - encoder (IntegerEncoder): An index to vocabulary encoder for the generator. - vocab_size (int): The size of the vocabulary for both discriminator and generator. - embedding_size (int): The size of the embedding layer for both discriminator and generator. - hidden_size (int): The amount of hidden neurons for the generator. - tau (float): Gumbel-Softmax temperature parameter. + encoder: An index to vocabulary encoder for the generator. + vocab_size: The size of the vocabulary for both discriminator and generator. + embedding_size: The size of the embedding layer for both discriminator and generator. + hidden_size: The amount of hidden neurons for the generator. + tau: Gumbel-Softmax temperature parameter. """ @@ -56,32 +65,37 @@ def __init__( logger.info("Class overrided.") @property - def vocab_size(self): - """int: The size of the vocabulary.""" + def vocab_size(self) -> int: + """The size of the vocabulary.""" return self._vocab_size @vocab_size.setter - def vocab_size(self, vocab_size): + def vocab_size(self, vocab_size: int) -> None: self._vocab_size = vocab_size @property - def init_tau(self): - """float: Gumbel-Softmax initial temperature.""" + def init_tau(self) -> float: + """Gumbel-Softmax initial temperature.""" return self._init_tau @init_tau.setter - def init_tau(self, init_tau): + def init_tau(self, init_tau: float) -> None: self._init_tau = init_tau - def compile(self, pre_optimizer, d_optimizer, g_optimizer): + def compile( + self, + pre_optimizer: tf.keras.optimizers, + d_optimizer: tf.keras.optimizers, + g_optimizer: tf.keras.optimizers, + ) -> None: """Main building method. Args: - pre_optimizer (tf.keras.optimizers): An optimizer instance for pre-training the generator. - d_optimizer (tf.keras.optimizers): An optimizer instance for the discriminator. - g_optimizer (tf.keras.optimizers): An optimizer instance for the generator. + pre_optimizer: An optimizer instance for pre-training the generator. + d_optimizer: An optimizer instance for the discriminator. + g_optimizer: An optimizer instance for the generator. """ @@ -102,7 +116,7 @@ def compile(self, pre_optimizer, d_optimizer, g_optimizer): self.history["D_loss"] = [] self.history["G_loss"] = [] - def generate_batch(self, x): + def generate_batch(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]: """Generates a batch of tokens by feeding to the network the current token (t) and predicting the next token (t+1). @@ -110,7 +124,7 @@ def generate_batch(self, x): x (tf.tensor): A tensor containing the inputs. Returns: - A (batch_size, length) tensor of generated tokens and a + (Tuple[tf.Tensor, tf.Tensor]): A (batch_size, length) tensor of generated tokens and a (batch_size, length, vocab_size) tensor of predictions. """ @@ -146,15 +160,15 @@ def generate_batch(self, x): return sampled_batch, sampled_preds - def _discriminator_loss(self, y_real, y_fake): + def _discriminator_loss(self, y_real: tf.Tensor, y_fake: tf.Tensor) -> tf.Tensor: """Calculates the loss out of the discriminator architecture. Args: - y_real (tf.tensor): A tensor containing the real data targets. - y_fake (tf.tensor): A tensor containing the fake data targets. + y_real: A tensor containing the real data targets. + y_fake: A tensor containing the fake data targets. Returns: - The loss based on the discriminator network. + (tf.Tensor): The loss based on the discriminator network. """ @@ -163,14 +177,14 @@ def _discriminator_loss(self, y_real, y_fake): return tf.reduce_mean(real_loss) + tf.reduce_mean(fake_loss) - def _generator_loss(self, y_fake): + def _generator_loss(self, y_fake: tf.Tensor) -> tf.Tensor: """Calculates the loss out of the generator architecture. Args: - y_fake (tf.tensor): A tensor containing the fake data targets. + y_fake: A tensor containing the fake data targets. Returns: - The loss based on the generator network. + (tf.Tensor): The loss based on the generator network. """ @@ -179,12 +193,12 @@ def _generator_loss(self, y_fake): return tf.reduce_mean(loss) @tf.function - def G_pre_step(self, x, y): + def G_pre_step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization pre-fitting step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -208,12 +222,12 @@ def G_pre_step(self, x, y): self.G_loss.update_state(loss) @tf.function - def step(self, x, y): + def step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization step. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -246,12 +260,12 @@ def step(self, x, y): self.D_loss.update_state(D_loss) self.G_loss.update_state(G_loss) - def pre_fit(self, batches, epochs=100): + def pre_fit(self, batches: Dataset, epochs: Optional[int] = 100) -> None: """Pre-trains the model. Args: - batches (Dataset): Pre-training batches containing samples. - epochs (int): The maximum number of pre-training epochs. + batches: Pre-training batches containing samples. + epochs: The maximum number of pre-training epochs. """ @@ -281,12 +295,12 @@ def pre_fit(self, batches, epochs=100): logger.to_file("Loss(G): %s", self.G_loss.result().numpy()) - def fit(self, batches, epochs=100): + def fit(self, batches: Dataset, epochs: Optional[int] = 100) -> None: """Trains the model. Args: - batches (Dataset): Training batches containing samples. - epochs (int): The maximum number of training epochs. + batches: Training batches containing samples. + epochs: The maximum number of training epochs. """ diff --git a/nalp/models/maligan.py b/nalp/models/maligan.py index 8c59f07..6452302 100644 --- a/nalp/models/maligan.py +++ b/nalp/models/maligan.py @@ -1,12 +1,16 @@ """Maximum-Likelihood Augmented Discrete Generative Adversarial Network. """ +from typing import Optional, Tuple + import numpy as np import tensorflow as tf from tensorflow.keras.utils import Progbar import nalp.utils.constants as c from nalp.core import Adversarial +from nalp.core.dataset import Dataset +from nalp.encoders.integer import IntegerEncoder from nalp.models.discriminators import EmbeddedTextDiscriminator from nalp.models.generators import LSTMGenerator from nalp.utils import logging @@ -26,28 +30,28 @@ class MaliGAN(Adversarial): def __init__( self, - encoder=None, - vocab_size=1, - max_length=1, - embedding_size=32, - hidden_size=64, - n_filters=(64), - filters_size=(1), - dropout_rate=0.25, - temperature=1, - ): + encoder: Optional[IntegerEncoder] = None, + vocab_size: Optional[int] = 1, + max_length: Optional[int] = 1, + embedding_size: Optional[int] = 32, + hidden_size: Optional[int] = 64, + n_filters: Optional[Tuple[int, ...]] = (64), + filters_size: Optional[Tuple[int, ...]] = (1), + dropout_rate: Optional[float] = 0.25, + temperature: Optional[float] = 1.0, + ) -> None: """Initialization method. Args: - encoder (IntegerEncoder): An index to vocabulary encoder for the generator. - vocab_size (int): The size of the vocabulary for both discriminator and generator. - max_length (int): Maximum length of the sequences for the discriminator. - embedding_size (int): The size of the embedding layer for both discriminator and generator. - hidden_size (int): The amount of hidden neurons for the generator. - n_filters (tuple): Number of filters to be applied in the discriminator. - filters_size (tuple): Size of filters to be applied in the discriminator. - dropout_rate (float): Dropout activation rate. - temperature (float): Temperature value to sample the token. + encoder: An index to vocabulary encoder for the generator. + vocab_size: The size of the vocabulary for both discriminator and generator. + max_length: Maximum length of the sequences for the discriminator. + embedding_size: The size of the embedding layer for both discriminator and generator. + hidden_size: The amount of hidden neurons for the generator. + n_filters: Number of filters to be applied in the discriminator. + filters_size: Size of filters to be applied in the discriminator. + dropout_rate: Dropout activation rate. + temperature: Temperature value to sample the token. """ @@ -77,32 +81,37 @@ def __init__( logger.info("Class overrided.") @property - def vocab_size(self): - """int: The size of the vocabulary.""" + def vocab_size(self) -> int: + """The size of the vocabulary.""" return self._vocab_size @vocab_size.setter - def vocab_size(self, vocab_size): + def vocab_size(self, vocab_size: int) -> None: self._vocab_size = vocab_size @property - def T(self): - """float: Temperature value to sample the token.""" + def T(self) -> float: + """Temperature value to sample the token.""" return self._T @T.setter - def T(self, T): + def T(self, T: float) -> None: self._T = T - def compile(self, pre_optimizer, d_optimizer, g_optimizer): + def compile( + self, + pre_optimizer: tf.keras.optimizers, + d_optimizer: tf.keras.optimizers, + g_optimizer: tf.keras.optimizers, + ) -> None: """Main building method. Args: - pre_optimizer (tf.keras.optimizers): An optimizer instance for pre-training the generator. - d_optimizer (tf.keras.optimizers): An optimizer instance for the discriminator. - g_optimizer (tf.keras.optimizers): An optimizer instance for the generator. + pre_optimizer: An optimizer instance for pre-training the generator. + d_optimizer: An optimizer instance for the discriminator. + g_optimizer: An optimizer instance for the generator. """ @@ -124,17 +133,18 @@ def compile(self, pre_optimizer, d_optimizer, g_optimizer): self.history["D_loss"] = [] self.history["G_loss"] = [] - def generate_batch(self, batch_size=1, length=1): + def generate_batch( + self, batch_size: Optional[int] = 1, length: Optional[int] = 1 + ) -> tf.Tensor: """Generates a batch of tokens by feeding to the network the current token (t) and predicting the next token (t+1). Args: - batch_size (int): Size of the batch to be generated. - length (int): Length of generated tokens. - temperature (float): A temperature value to sample the token. + batch_size: Size of the batch to be generated. + length: Length of generated tokens. Returns: - A (batch_size, length) tensor of generated tokens. + (tf.Tensor): A (batch_size, length) tensor of generated tokens. """ @@ -174,11 +184,14 @@ def generate_batch(self, batch_size=1, length=1): return x_sampled_batch, y_sampled_batch - def _get_reward(self, x): + def _get_reward(self, x: tf.Tensor) -> tf.Tensor: """Calculates rewards over an input using a Maximum-Likelihood approach. Args: - x (tf.tensor): A tensor containing the inputs. + x: A tensor containing the inputs. + + Returns: + (tf.Tensor): Reward over input. """ @@ -200,12 +213,12 @@ def _get_reward(self, x): return rewards @tf.function - def G_pre_step(self, x, y): + def G_pre_step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization pre-fitting step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -227,13 +240,13 @@ def G_pre_step(self, x, y): self.G_loss.update_state(loss) @tf.function - def G_step(self, x, y, rewards): + def G_step(self, x: tf.Tensor, y: tf.Tensor, rewards: tf.Tensor) -> None: """Performs a single batch optimization step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. - rewards (tf.tensor): A tensor containing the rewards for the input. + x : A tensor containing the inputs. + y: A tensor containing the inputs' labels. + rewards: A tensor containing the rewards for the input. """ @@ -255,12 +268,12 @@ def G_step(self, x, y, rewards): self.G_loss.update_state(loss) @tf.function - def D_step(self, x, y): + def D_step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization step over the discriminator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -281,13 +294,18 @@ def D_step(self, x, y): # Updates the discriminator's loss state self.D_loss.update_state(loss) - def pre_fit(self, batches, g_epochs=50, d_epochs=10): + def pre_fit( + self, + batches: Dataset, + g_epochs: Optional[int] = 50, + d_epochs: Optional[int] = 10, + ) -> None: """Pre-trains the model. Args: - batches (Dataset): Pre-training batches containing samples. - g_epochs (int): The maximum number of pre-training generator epochs. - d_epochs (int): The maximum number of pre-training discriminator epochs. + batches: Pre-training batches containing samples. + g_epochs: The maximum number of pre-training generator epochs. + d_epochs: The maximum number of pre-training discriminator epochs. """ @@ -370,13 +388,15 @@ def pre_fit(self, batches, g_epochs=50, d_epochs=10): logger.to_file("Loss(D): %s", self.D_loss.result().numpy()) - def fit(self, batches, epochs=10, d_epochs=5): + def fit( + self, batches: Dataset, epochs: Optional[int] = 10, d_epochs: Optional[int] = 5 + ) -> None: """Trains the model. Args: - batches (Dataset): Training batches containing samples. - epochs (int): The maximum number of total training epochs. - d_epochs (int): The maximum number of discriminator epochs per total epoch. + batches: Training batches containing samples. + epochs: The maximum number of total training epochs. + d_epochs: The maximum number of discriminator epochs per total epoch. """ diff --git a/nalp/models/relgan.py b/nalp/models/relgan.py index 48ee261..6cdaecf 100644 --- a/nalp/models/relgan.py +++ b/nalp/models/relgan.py @@ -1,10 +1,14 @@ """Relational Generative Adversarial Network. """ +from typing import Optional, Tuple + import tensorflow as tf from tensorflow.keras.utils import Progbar from nalp.core import Adversarial +from nalp.core.dataset import Dataset +from nalp.encoders.integer import IntegerEncoder from nalp.models.discriminators import TextDiscriminator from nalp.models.generators import GumbelRMCGenerator from nalp.utils import logging @@ -23,36 +27,36 @@ class RelGAN(Adversarial): def __init__( self, - encoder=None, - vocab_size=1, - max_length=1, - embedding_size=32, - n_slots=3, - n_heads=5, - head_size=10, - n_blocks=1, - n_layers=3, - n_filters=(64), - filters_size=(1), - dropout_rate=0.25, - tau=5, + encoder: Optional[IntegerEncoder] = None, + vocab_size: Optional[int] = 1, + max_length: Optional[int] = 1, + embedding_size: Optional[int] = 32, + n_slots: Optional[int] = 3, + n_heads: Optional[int] = 5, + head_size: Optional[int] = 10, + n_blocks: Optional[int] = 1, + n_layers: Optional[int] = 3, + n_filters: Optional[Tuple[int, ...]] = (64), + filters_size: Optional[Tuple[int, ...]] = (1), + dropout_rate: Optional[float] = 0.25, + tau: Optional[float] = 5.0, ): """Initialization method. Args: - encoder (IntegerEncoder): An index to vocabulary encoder for the generator. - vocab_size (int): The size of the vocabulary for both discriminator and generator. - max_length (int): Maximum length of the sequences for the discriminator. - embedding_size (int): The size of the embedding layer for both discriminator and generator. - n_slots (int): Number of memory slots for the generator. - n_heads (int): Number of attention heads for the generator. - head_size (int): Size of each attention head for the generator. - n_blocks (int): Number of feed-forward networks for the generator. - n_layers (int): Amout of layers per feed-forward network for the generator. - n_filters (tuple): Number of filters to be applied in the discriminator. - filters_size (tuple): Size of filters to be applied in the discriminator. - dropout_rate (float): Dropout activation rate. - tau (float): Gumbel-Softmax temperature parameter. + encoder: An index to vocabulary encoder for the generator. + vocab_size: The size of the vocabulary for both discriminator and generator. + max_length: Maximum length of the sequences for the discriminator. + embedding_size: The size of the embedding layer for both discriminator and generator. + n_slots: Number of memory slots for the generator. + n_heads: Number of attention heads for the generator. + head_size: Size of each attention head for the generator. + n_blocks: Number of feed-forward networks for the generator. + n_layers: Amout of layers per feed-forward network for the generator. + n_filters: Number of filters to be applied in the discriminator. + filters_size: Size of filters to be applied in the discriminator. + dropout_rate: Dropout activation rate. + tau: Gumbel-Softmax temperature parameter. """ @@ -87,32 +91,37 @@ def __init__( logger.info("Class overrided.") @property - def vocab_size(self): - """int: The size of the vocabulary.""" + def vocab_size(self) -> int: + """The size of the vocabulary.""" return self._vocab_size @vocab_size.setter - def vocab_size(self, vocab_size): + def vocab_size(self, vocab_size: int) -> None: self._vocab_size = vocab_size @property - def init_tau(self): - """float: Gumbel-Softmax initial temperature.""" + def init_tau(self) -> float: + """Gumbel-Softmax initial temperature.""" return self._init_tau @init_tau.setter - def init_tau(self, init_tau): + def init_tau(self, init_tau: float) -> None: self._init_tau = init_tau - def compile(self, pre_optimizer, d_optimizer, g_optimizer): + def compile( + self, + pre_optimizer: tf.keras.optimizers, + d_optimizer: tf.keras.optimizers, + g_optimizer: tf.keras.optimizers, + ) -> None: """Main building method. Args: - pre_optimizer (tf.keras.optimizers): An optimizer instance for pre-training the generator. - d_optimizer (tf.keras.optimizers): An optimizer instance for the discriminator. - g_optimizer (tf.keras.optimizers): An optimizer instance for the generator. + pre_optimizer: An optimizer instance for pre-training the generator. + d_optimizer: An optimizer instance for the discriminator. + g_optimizer: An optimizer instance for the generator. """ @@ -133,15 +142,15 @@ def compile(self, pre_optimizer, d_optimizer, g_optimizer): self.history["D_loss"] = [] self.history["G_loss"] = [] - def generate_batch(self, x): + def generate_batch(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]: """Generates a batch of tokens by feeding to the network the current token (t) and predicting the next token (t+1). Args: - x (tf.tensor): A tensor containing the inputs. + x: A tensor containing the inputs. Returns: - A (batch_size, length) tensor of generated tokens and a + (Tuple[tf.Tensor, tf.Tensor]): A (batch_size, length) tensor of generated tokens and a (batch_size, length, vocab_size) tensor of predictions. """ @@ -177,15 +186,15 @@ def generate_batch(self, x): return sampled_batch, sampled_preds - def _discriminator_loss(self, y_real, y_fake): + def _discriminator_loss(self, y_real: tf.Tensor, y_fake: tf.Tensor) -> tf.Tensor: """Calculates the loss out of the discriminator architecture. Args: - y_real (tf.tensor): A tensor containing the real data targets. - y_fake (tf.tensor): A tensor containing the fake data targets. + y_real: A tensor containing the real data targets. + y_fake: A tensor containing the fake data targets. Returns: - The loss based on the discriminator network. + (tf.Tensor): The loss based on the discriminator network. """ @@ -193,15 +202,15 @@ def _discriminator_loss(self, y_real, y_fake): return tf.reduce_mean(loss) - def _generator_loss(self, y_real, y_fake): + def _generator_loss(self, y_real: tf.Tensor, y_fake: tf.Tensor) -> tf.Tensor: """Calculates the loss out of the generator architecture. Args: - y_real (tf.tensor): A tensor containing the real data targets. - y_fake (tf.tensor): A tensor containing the fake data targets. + y_real: A tensor containing the real data targets. + y_fake: A tensor containing the fake data targets. Returns: - The loss based on the generator network. + (tf.Tensor): The loss based on the generator network. """ @@ -210,12 +219,12 @@ def _generator_loss(self, y_real, y_fake): return tf.reduce_mean(loss) @tf.function - def G_pre_step(self, x, y): + def G_pre_step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization pre-fitting step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -239,12 +248,12 @@ def G_pre_step(self, x, y): self.G_loss.update_state(loss) @tf.function - def step(self, x, y): + def step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization step. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -277,12 +286,12 @@ def step(self, x, y): self.G_loss.update_state(G_loss) self.D_loss.update_state(D_loss) - def pre_fit(self, batches, epochs=100): + def pre_fit(self, batches: Dataset, epochs: Optional[int] = 100) -> None: """Pre-trains the model. Args: - batches (Dataset): Pre-training batches containing samples. - epochs (int): The maximum number of pre-training epochs. + batches: Pre-training batches containing samples. + epochs: The maximum number of pre-training epochs. """ @@ -312,12 +321,12 @@ def pre_fit(self, batches, epochs=100): logger.to_file("Loss(G): %s", self.G_loss.result().numpy()) - def fit(self, batches, epochs=100): + def fit(self, batches: Dataset, epochs: Optional[int] = 100) -> None: """Trains the model. Args: - batches (Dataset): Training batches containing samples. - epochs (int): The maximum number of training epochs. + batches: Training batches containing samples. + epochs: The maximum number of training epochs. """ diff --git a/nalp/models/seqgan.py b/nalp/models/seqgan.py index f2a4f53..a72b966 100644 --- a/nalp/models/seqgan.py +++ b/nalp/models/seqgan.py @@ -1,12 +1,16 @@ """Sequence Generative Adversarial Network. """ +from typing import Optional, Tuple + import numpy as np import tensorflow as tf from tensorflow.keras.utils import Progbar import nalp.utils.constants as c from nalp.core import Adversarial +from nalp.core.dataset import Dataset +from nalp.encoders.integer import IntegerEncoder from nalp.models.discriminators import EmbeddedTextDiscriminator from nalp.models.generators import LSTMGenerator from nalp.utils import logging @@ -25,28 +29,28 @@ class SeqGAN(Adversarial): def __init__( self, - encoder=None, - vocab_size=1, - max_length=1, - embedding_size=32, - hidden_size=64, - n_filters=(64), - filters_size=(1), - dropout_rate=0.25, - temperature=1, - ): + encoder: Optional[IntegerEncoder] = None, + vocab_size: Optional[int] = 1, + max_length: Optional[int] = 1, + embedding_size: Optional[int] = 32, + hidden_size: Optional[int] = 64, + n_filters: Optional[Tuple[int, ...]] = (64), + filters_size: Optional[Tuple[int, ...]] = (1), + dropout_rate: Optional[float] = 0.25, + temperature: Optional[float] = 1.0, + ) -> None: """Initialization method. Args: - encoder (IntegerEncoder): An index to vocabulary encoder for the generator. - vocab_size (int): The size of the vocabulary for both discriminator and generator. - max_length (int): Maximum length of the sequences for the discriminator. - embedding_size (int): The size of the embedding layer for both discriminator and generator. - hidden_size (int): The amount of hidden neurons for the generator. - n_filters (tuple): Number of filters to be applied in the discriminator. - filters_size (tuple): Size of filters to be applied in the discriminator. - dropout_rate (float): Dropout activation rate. - temperature (float): Temperature value to sample the token. + encoder: An index to vocabulary encoder for the generator. + vocab_size: The size of the vocabulary for both discriminator and generator. + max_length: Maximum length of the sequences for the discriminator. + embedding_size: The size of the embedding layer for both discriminator and generator. + hidden_size: The amount of hidden neurons for the generator. + n_filters: Number of filters to be applied in the discriminator. + filters_size: Size of filters to be applied in the discriminator. + dropout_rate: Dropout activation rate. + temperature: Temperature value to sample the token. """ @@ -76,32 +80,37 @@ def __init__( logger.info("Class overrided.") @property - def vocab_size(self): - """int: The size of the vocabulary.""" + def vocab_size(self) -> int: + """The size of the vocabulary.""" return self._vocab_size @vocab_size.setter - def vocab_size(self, vocab_size): + def vocab_size(self, vocab_size: int) -> None: self._vocab_size = vocab_size @property - def T(self): - """float: Temperature value to sample the token.""" + def T(self) -> float: + """Temperature value to sample the token.""" return self._T @T.setter - def T(self, T): + def T(self, T: float) -> None: self._T = T - def compile(self, pre_optimizer, d_optimizer, g_optimizer): + def compile( + self, + pre_optimizer: tf.keras.optimizers, + d_optimizer: tf.keras.optimizers, + g_optimizer: tf.keras.optimizers, + ) -> None: """Main building method. Args: - pre_optimizer (tf.keras.optimizers): An optimizer instance for pre-training the generator. - d_optimizer (tf.keras.optimizers): An optimizer instance for the discriminator. - g_optimizer (tf.keras.optimizers): An optimizer instance for the generator. + pre_optimizer: An optimizer instance for pre-training the generator. + d_optimizer: An optimizer instance for the discriminator. + g_optimizer: An optimizer instance for the generator. """ @@ -123,17 +132,18 @@ def compile(self, pre_optimizer, d_optimizer, g_optimizer): self.history["D_loss"] = [] self.history["G_loss"] = [] - def generate_batch(self, batch_size=1, length=1): + def generate_batch( + self, batch_size: Optional[int] = 1, length: Optional[int] = 1 + ) -> tf.Tensor: """Generates a batch of tokens by feeding to the network the current token (t) and predicting the next token (t+1). Args: - batch_size (int): Size of the batch to be generated. - length (int): Length of generated tokens. - temperature (float): A temperature value to sample the token. + batch_size: Size of the batch to be generated. + length: Length of generated tokens. Returns: - A (batch_size, length) tensor of generated tokens. + (tf.Tensor): A (batch_size, length) tensor of generated tokens. """ @@ -173,12 +183,15 @@ def generate_batch(self, batch_size=1, length=1): return x_sampled_batch, y_sampled_batch - def _get_reward(self, x, n_rollouts): + def _get_reward(self, x: tf.Tensor, n_rollouts: int) -> tf.Tensor: """Calculates rewards over an input using a Monte Carlo search strategy. Args: - x (tf.tensor): A tensor containing the inputs. - n_rollouts (int): Number of rollouts for conducting the Monte Carlo search. + x: A tensor containing the inputs. + n_rollouts: Number of rollouts for conducting the Monte Carlo search. + + Returns: + (tf.Tensor): Reward over input. """ @@ -225,12 +238,12 @@ def _get_reward(self, x, n_rollouts): return rewards @tf.function - def G_pre_step(self, x, y): + def G_pre_step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization pre-fitting step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -252,13 +265,13 @@ def G_pre_step(self, x, y): self.G_loss.update_state(loss) @tf.function - def G_step(self, x, y, rewards): + def G_step(self, x: tf.Tensor, y: tf.Tensor, rewards: tf.Tensor) -> None: """Performs a single batch optimization step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. - rewards (tf.tensor): A tensor containing the rewards for the input. + x : A tensor containing the inputs. + y: A tensor containing the inputs' labels. + rewards: A tensor containing the rewards for the input. """ @@ -280,12 +293,12 @@ def G_step(self, x, y, rewards): self.G_loss.update_state(loss) @tf.function - def D_step(self, x, y): + def D_step(self, x: tf.Tensor, y: tf.Tensor) -> None: """Performs a single batch optimization step over the discriminator. Args: - x (tf.tensor): A tensor containing the inputs. - y (tf.tensor): A tensor containing the inputs' labels. + x: A tensor containing the inputs. + y: A tensor containing the inputs' labels. """ @@ -306,13 +319,18 @@ def D_step(self, x, y): # Updates the discriminator's loss state self.D_loss.update_state(loss) - def pre_fit(self, batches, g_epochs=50, d_epochs=10): + def pre_fit( + self, + batches: Dataset, + g_epochs: Optional[int] = 50, + d_epochs: Optional[int] = 10, + ) -> None: """Pre-trains the model. Args: - batches (Dataset): Pre-training batches containing samples. - g_epochs (int): The maximum number of pre-training generator epochs. - d_epochs (int): The maximum number of pre-training discriminator epochs. + batches: Pre-training batches containing samples. + g_epochs: The maximum number of pre-training generator epochs. + d_epochs: The maximum number of pre-training discriminator epochs. """ @@ -395,15 +413,22 @@ def pre_fit(self, batches, g_epochs=50, d_epochs=10): logger.to_file("Loss(D): %s", self.D_loss.result().numpy()) - def fit(self, batches, epochs=10, g_epochs=1, d_epochs=5, n_rollouts=16): + def fit( + self, + batches: Dataset, + epochs: Optional[int] = 10, + g_epochs: Optional[int] = 1, + d_epochs: Optional[int] = 5, + n_rollouts: Optional[int] = 16, + ) -> None: """Trains the model. Args: - batches (Dataset): Training batches containing samples. - epochs (int): The maximum number of total training epochs. - g_epochs (int): The maximum number of generator epochs per total epoch. - d_epochs (int): The maximum number of discriminator epochs per total epoch. - n_rollouts (int): Number of rollouts for conducting the Monte Carlo search. + batches: Training batches containing samples. + epochs: The maximum number of total training epochs. + g_epochs: The maximum number of generator epochs per total epoch. + d_epochs: The maximum number of discriminator epochs per total epoch. + n_rollouts: Number of rollouts for conducting the Monte Carlo search. """ diff --git a/nalp/models/wgan.py b/nalp/models/wgan.py index c819995..7874e3e 100644 --- a/nalp/models/wgan.py +++ b/nalp/models/wgan.py @@ -1,10 +1,13 @@ """Wasserstein Generative Adversarial Network. """ +from typing import Optional, Tuple + import tensorflow as tf from tensorflow.keras.utils import Progbar from nalp.core import Adversarial +from nalp.core.dataset import Dataset from nalp.models.discriminators import ConvDiscriminator from nalp.models.generators import ConvGenerator from nalp.utils import logging @@ -29,26 +32,26 @@ class WGAN(Adversarial): def __init__( self, - input_shape=(28, 28, 1), - noise_dim=100, - n_samplings=3, - alpha=0.3, - dropout_rate=0.3, - model_type="wc", - clip=0.01, - penalty=10, + input_shape: Optional[Tuple[int, int, int]] = (28, 28, 1), + noise_dim: Optional[int] = 100, + n_samplings: Optional[int] = 3, + alpha: Optional[float] = 0.3, + dropout_rate: Optional[float] = 0.3, + model_type: Optional[str] = "wc", + clip: Optional[float] = 0.01, + penalty: Optional[int] = 10, ): """Initialization method. Args: - input_shape (tuple): An input shape for the Generator. - noise_dim (int): Amount of noise dimensions for the Generator. - n_samplings (int): Number of down/up samplings to perform. - alpha (float): LeakyReLU activation threshold. - dropout_rate (float): Dropout activation rate. - model_type (str): Whether should use weight clipping (wc) or gradient penalty (gp). - clip (float): Clipping value for the Lipschitz constrain. - penalty (int): Coefficient for the gradient penalty. + input_shape: An input shape for the Generator. + noise_dim: Amount of noise dimensions for the Generator. + n_samplings: Number of down/up samplings to perform. + alpha: LeakyReLU activation threshold. + dropout_rate: Dropout activation rate. + model_type: Whether should use weight clipping (wc) or gradient penalty (gp). + clip: Clipping value for the Lipschitz constrain. + penalty: Coefficient for the gradient penalty. """ @@ -85,44 +88,44 @@ def __init__( logger.info("Class overrided.") @property - def model_type(self): - """str: Whether should use weight clipping (wc) or gradient penalty (gp).""" + def model_type(self) -> str: + """Whether should use weight clipping (wc) or gradient penalty (gp).""" return self._model_type @model_type.setter - def model_type(self, model_type): + def model_type(self, model_type: str) -> None: self._model_type = model_type @property - def clip(self): - """float: Clipping value for the Lipschitz constrain.""" + def clip(self) -> float: + """Clipping value for the Lipschitz constrain.""" return self._clip @clip.setter - def clip(self, clip): + def clip(self, clip: float) -> None: self._clip = clip @property - def penalty_lambda(self): - """int: Coefficient for the gradient penalty.""" + def penalty_lambda(self) -> int: + """Coefficient for the gradient penalty.""" return self._penalty_lambda @penalty_lambda.setter - def penalty_lambda(self, penalty_lambda): + def penalty_lambda(self, penalty_lambda: int) -> None: self._penalty_lambda = penalty_lambda - def _gradient_penalty(self, x, x_fake): + def _gradient_penalty(self, x: tf.Tensor, x_fake: tf.Tensor) -> tf.Tensor: """Performs the gradient penalty procedure. Args: - x (tf.tensor): A tensor containing the real inputs. - x_fake (tf.tensor): A tensor containing the fake inputs. + x: A tensor containing the real inputs. + x_fake: A tensor containing the fake inputs. Returns: - The penalization to be applied over the loss function. + (tf.Tensor): The penalization to be applied over the loss function. """ @@ -154,11 +157,11 @@ def _gradient_penalty(self, x, x_fake): return penalty @tf.function - def D_step(self, x): + def D_step(self, x: tf.Tensor) -> None: """Performs a single batch optimization step over the discriminator. Args: - x (tf.tensor): A tensor containing the inputs. + x: A tensor containing the inputs. """ @@ -203,11 +206,11 @@ def D_step(self, x): ] @tf.function - def G_step(self, x): + def G_step(self, x: tf.Tensor) -> None: """Performs a single batch optimization step over the generator. Args: - x (tf.tensor): A tensor containing the inputs. + x: A tensor containing the inputs. """ @@ -234,13 +237,18 @@ def G_step(self, x): # Updates the generator's loss state self.G_loss.update_state(G_loss) - def fit(self, batches, epochs=100, critic_steps=5): + def fit( + self, + batches: Dataset, + epochs: Optional[int] = 100, + critic_steps: Optional[int] = 5, + ) -> None: """Trains the model. Args: - batches (Dataset): Training batches containing samples. - epochs (int): The maximum number of training epochs. - critic_steps (int): Amount of discriminator epochs per training epoch. + batches: Training batches containing samples. + epochs: The maximum number of training epochs. + critic_steps: Amount of discriminator epochs per training epoch. """