diff --git a/nalp/models/discriminators/conv.py b/nalp/models/discriminators/conv.py
index 0b064d1..6477369 100644
--- a/nalp/models/discriminators/conv.py
+++ b/nalp/models/discriminators/conv.py
@@ -1,6 +1,8 @@
 """Convolutional discriminator.
 """
 
+from typing import Optional
+
 import tensorflow as tf
 from tensorflow.keras.layers import Conv2D, Dense, Dropout
 
@@ -16,13 +18,18 @@ class ConvDiscriminator(Discriminator):
 
     """
 
-    def __init__(self, n_samplings=3, alpha=0.3, dropout_rate=0.3):
+    def __init__(
+        self,
+        n_samplings: Optional[int] = 3,
+        alpha: Optional[float] = 0.3,
+        dropout_rate: Optional[float] = 0.3,
+    ) -> None:
         """Initialization method.
 
         Args:
-            n_samplings (int): Number of downsamplings to perform.
-            alpha (float): LeakyReLU activation threshold.
-            dropout_rate (float): Dropout activation rate.
+            n_samplings: Number of downsamplings to perform.
+            alpha: LeakyReLU activation threshold.
+            dropout_rate: Dropout activation rate.
 
         """
 
@@ -52,24 +59,24 @@ def __init__(self, n_samplings=3, alpha=0.3, dropout_rate=0.3):
         logger.info("Class overrided.")
 
     @property
-    def alpha(self):
-        """float: LeakyReLU activation threshold."""
+    def alpha(self) -> float:
+        """LeakyReLU activation threshold."""
 
         return self._alpha
 
     @alpha.setter
-    def alpha(self, alpha):
+    def alpha(self, alpha: float) -> None:
         self._alpha = alpha
 
-    def call(self, x, training=True):
+    def call(self, x: tf.Tensor, training: Optional[bool] = True) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            training (bool): Whether architecture is under training or not.
+            x: A tensorflow's tensor holding input data.
+            training: Whether architecture is under training or not.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/discriminators/embedded_text.py b/nalp/models/discriminators/embedded_text.py
index 7f85653..2d080df 100644
--- a/nalp/models/discriminators/embedded_text.py
+++ b/nalp/models/discriminators/embedded_text.py
@@ -1,6 +1,8 @@
 """Embedded-text discriminator.
 """
 
+from typing import Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import Conv2D, Dense, Dropout, Embedding, MaxPool1D
 
@@ -18,22 +20,22 @@ class EmbeddedTextDiscriminator(Discriminator):
 
     def __init__(
         self,
-        vocab_size=1,
-        max_length=1,
-        embedding_size=32,
-        n_filters=(64),
-        filters_size=(1),
-        dropout_rate=0.25,
-    ):
+        vocab_size: Optional[int] = 1,
+        max_length: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        n_filters: Optional[Tuple[int, ...]] = (64),
+        filters_size: Optional[Tuple[int, ...]] = (1),
+        dropout_rate: Optional[float] = 0.25,
+    ) -> None:
         """Initialization method.
 
         Args:
-            vocab_size (int): The size of the vocabulary.
-            max_length (int): Maximum length of the sequences.
-            embedding_size (int): The size of the embedding layer.
-            n_filters (tuple): Number of filters to be applied.
-            filters_size (tuple): Size of filters to be applied.
-            dropout_rate (float): Dropout activation rate.
+            vocab_size: The size of the vocabulary.
+            max_length: Maximum length of the sequences.
+            embedding_size: The size of the embedding layer.
+            n_filters: Number of filters to be applied.
+            filters_size: Size of filters to be applied.
+            dropout_rate: Dropout activation rate.
 
         """
 
@@ -72,15 +74,15 @@ def __init__(
 
         logger.info("Class overrided.")
 
-    def call(self, x, training=True):
+    def call(self, x: tf.Tensor, training: Optional[bool] = True) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            training (bool): Whether architecture is under training or not.
+            x: A tensorflow's tensor holding input data.
+            training: Whether architecture is under training or not.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/discriminators/linear.py b/nalp/models/discriminators/linear.py
index a9bf4e4..59af570 100644
--- a/nalp/models/discriminators/linear.py
+++ b/nalp/models/discriminators/linear.py
@@ -1,6 +1,8 @@
 """Linear discriminator.
 """
 
+from typing import Optional
+
 import tensorflow as tf
 from tensorflow.keras.layers import Dense
 
@@ -16,12 +18,14 @@ class LinearDiscriminator(Discriminator):
 
     """
 
-    def __init__(self, n_samplings=3, alpha=0.01):
+    def __init__(
+        self, n_samplings: Optional[int] = 3, alpha: Optional[float] = 0.01
+    ) -> None:
         """Initialization method.
 
         Args:
-            n_samplings (int): Number of downsamplings to perform.
-            alpha (float): LeakyReLU activation threshold.
+            n_samplings: Number of downsamplings to perform.
+            alpha: LeakyReLU activation threshold.
 
         """
 
@@ -43,24 +47,24 @@ def __init__(self, n_samplings=3, alpha=0.01):
         logger.info("Class overrided.")
 
     @property
-    def alpha(self):
-        """float: LeakyReLU activation threshold."""
+    def alpha(self) -> float:
+        """LeakyReLU activation threshold."""
 
         return self._alpha
 
     @alpha.setter
-    def alpha(self, alpha):
+    def alpha(self, alpha: float) -> None:
         self._alpha = alpha
 
-    def call(self, x, training=True):
+    def call(self, x: tf.Tensor, training: Optional[bool] = True) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            training (bool): Whether architecture is under training or not.
+            x: A tensorflow's tensor holding input data.
+            training: Whether architecture is under training or not.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/discriminators/lstm.py b/nalp/models/discriminators/lstm.py
index 3c2b587..953ed38 100644
--- a/nalp/models/discriminators/lstm.py
+++ b/nalp/models/discriminators/lstm.py
@@ -1,6 +1,9 @@
 """Long Short-Term Memory discriminator.
 """
 
+from typing import Optional
+
+import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, LSTMCell
 
 from nalp.core import Discriminator
@@ -18,12 +21,14 @@ class LSTMDiscriminator(Discriminator):
 
     """
 
-    def __init__(self, embedding_size=32, hidden_size=64):
+    def __init__(
+        self, embedding_size: Optional[int] = 32, hidden_size: Optional[int] = 64
+    ) -> None:
         """Initialization method.
 
         Args:
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (int): The amount of hidden neurons.
+            embedding_size: The size of the embedding layer.
+            hidden_size: The amount of hidden neurons.
 
         """
 
@@ -47,14 +52,14 @@ def __init__(self, embedding_size=32, hidden_size=64):
 
         logger.info("Class overrided.")
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/discriminators/text.py b/nalp/models/discriminators/text.py
index eaac26b..8f79ebc 100644
--- a/nalp/models/discriminators/text.py
+++ b/nalp/models/discriminators/text.py
@@ -1,6 +1,8 @@
 """Text-based discriminator.
 """
 
+from typing import Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import Conv2D, Dense, Dropout, MaxPool1D
 
@@ -18,20 +20,20 @@ class TextDiscriminator(Discriminator):
 
     def __init__(
         self,
-        max_length=1,
-        embedding_size=32,
-        n_filters=(64),
-        filters_size=(1),
-        dropout_rate=0.25,
-    ):
+        max_length: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        n_filters: Optional[Tuple[int, ...]] = (64),
+        filters_size: Optional[Tuple[int, ...]] = (1),
+        dropout_rate: Optional[float] = 0.25,
+    ) -> None:
         """Initialization method.
 
         Args:
-            max_length (int): Maximum length of the sequences.
-            embedding_size (int): The size of the embedding layer.
-            n_filters (tuple): Number of filters to be applied.
-            filters_size (tuple): Size of filters to be applied.
-            dropout_rate (float): Dropout activation rate.
+            max_length: Maximum length of the sequences.
+            embedding_size: The size of the embedding layer.
+            n_filters: Number of filters to be applied.
+            filters_size: Size of filters to be applied.
+            dropout_rate: Dropout activation rate.
 
         """
 
@@ -67,15 +69,15 @@ def __init__(
 
         logger.info("Class overrided.")
 
-    def call(self, x, training=True):
+    def call(self, x: tf.Tensor, training: Optional[bool] = True) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            training (bool): Whether architecture is under training or not.
+            x: A tensorflow's tensor holding input data.
+            training: Whether architecture is under training or not.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/bi_lstm.py b/nalp/models/generators/bi_lstm.py
index 7d0d4b8..b60e2bc 100644
--- a/nalp/models/generators/bi_lstm.py
+++ b/nalp/models/generators/bi_lstm.py
@@ -1,10 +1,13 @@
 """Bi-directional Long Short-Term Memory generator.
 """
 
+from typing import Optional
+
 import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, Embedding, LSTMCell
 
 from nalp.core import Generator
+from nalp.encoders.integer import IntegerEncoder
 from nalp.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -19,14 +22,20 @@ class BiLSTMGenerator(Generator):
 
     """
 
-    def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64):
+    def __init__(
+        self,
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        hidden_size: Optional[int] = 64,
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (int): The amount of hidden neurons.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            hidden_size: The amount of hidden neurons.
 
         """
 
@@ -66,23 +75,23 @@ def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64
         logger.info("Class overrided.")
 
     @property
-    def encoder(self):
-        """obj: An encoder generic object."""
+    def encoder(self) -> IntegerEncoder:
+        """An encoder generic object."""
 
         return self._encoder
 
     @encoder.setter
-    def encoder(self, encoder):
+    def encoder(self, encoder: IntegerEncoder) -> None:
         self._encoder = encoder
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/conv.py b/nalp/models/generators/conv.py
index 235fe79..61eb278 100644
--- a/nalp/models/generators/conv.py
+++ b/nalp/models/generators/conv.py
@@ -1,6 +1,8 @@
 """Convolutional generator.
 """
 
+from typing import Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import BatchNormalization, Conv2DTranspose, Dense
 
@@ -17,15 +19,19 @@ class ConvGenerator(Generator):
     """
 
     def __init__(
-        self, input_shape=(28, 28, 1), noise_dim=100, n_samplings=3, alpha=0.3
-    ):
+        self,
+        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,
+    ) -> None:
         """Initialization method.
 
         Args:
-            input_shape (tuple): An input shape for the tensor.
-            noise_dim (int): Amount of noise dimensions.
-            n_samplings (int): Number of upsamplings to perform.
-            alpha (float): LeakyReLU activation threshold.
+            input_shape: An input shape for the tensor.
+            noise_dim: Amount of noise dimensions.
+            n_samplings: Number of upsamplings to perform.
+            alpha: LeakyReLU activation threshold.
 
         """
 
@@ -108,54 +114,54 @@ def __init__(
         logger.info("Class overrided.")
 
     @property
-    def alpha(self):
-        """float: LeakyReLU activation threshold."""
+    def alpha(self) -> float:
+        """LeakyReLU activation threshold."""
 
         return self._alpha
 
     @alpha.setter
-    def alpha(self, alpha):
+    def alpha(self, alpha: float) -> None:
         self._alpha = alpha
 
     @property
-    def noise_dim(self):
-        """int: Amount of noise dimensions."""
+    def noise_dim(self) -> int:
+        """Amount of noise dimensions."""
 
         return self._noise_dim
 
     @noise_dim.setter
-    def noise_dim(self, noise_dim):
+    def noise_dim(self, noise_dim: int) -> None:
         self._noise_dim = noise_dim
 
     @property
-    def sampling_factor(self):
-        """int: Sampling factor used to calculate the upsampling."""
+    def sampling_factor(self) -> int:
+        """Sampling factor used to calculate the upsampling."""
 
         return self._sampling_factor
 
     @sampling_factor.setter
-    def sampling_factor(self, sampling_factor):
+    def sampling_factor(self, sampling_factor: int) -> None:
         self._sampling_factor = sampling_factor
 
     @property
-    def filter_size(self):
-        """int: Initial size of the filter."""
+    def filter_size(self) -> int:
+        """Initial size of the filter."""
 
         return self._filter_size
 
     @filter_size.setter
-    def filter_size(self, filter_size):
+    def filter_size(self, filter_size: int) -> None:
         self._filter_size = filter_size
 
-    def call(self, x, training=True):
+    def call(self, x: tf.Tensor, training: Optional[bool] = True) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            training (bool): Whether architecture is under training or not.
+            x: A tensorflow's tensor holding input data.
+            training: Whether architecture is under training or not.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/gru.py b/nalp/models/generators/gru.py
index 59182b2..8c40eb8 100644
--- a/nalp/models/generators/gru.py
+++ b/nalp/models/generators/gru.py
@@ -1,9 +1,13 @@
 """Gated Recurrent Unit generator.
 """
 
+from typing import Optional
+
+import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, Embedding, GRUCell
 
 from nalp.core import Generator
+from nalp.encoders.integer import IntegerEncoder
 from nalp.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -20,14 +24,20 @@ class GRUGenerator(Generator):
 
     """
 
-    def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64):
+    def __init__(
+        self,
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        hidden_size: Optional[int] = 64,
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (int): The amount of hidden neurons.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            hidden_size: The amount of hidden neurons.
 
         """
 
@@ -55,23 +65,23 @@ def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64
         logger.info("Class overrided.")
 
     @property
-    def encoder(self):
-        """obj: An encoder generic object."""
+    def encoder(self) -> IntegerEncoder:
+        """An encoder generic object."""
 
         return self._encoder
 
     @encoder.setter
-    def encoder(self, encoder):
+    def encoder(self, encoder: IntegerEncoder) -> None:
         self._encoder = encoder
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/gumbel_lstm.py b/nalp/models/generators/gumbel_lstm.py
index fe0113f..344861c 100644
--- a/nalp/models/generators/gumbel_lstm.py
+++ b/nalp/models/generators/gumbel_lstm.py
@@ -1,9 +1,12 @@
 """Gumbel Long Short-Term Memory generator.
 """
 
+from typing import List, Optional, Tuple
+
 import tensorflow as tf
 
 import nalp.utils.constants as c
+from nalp.encoders.integer import IntegerEncoder
 from nalp.models.generators import LSTMGenerator
 from nalp.models.layers import GumbelSoftmax
 from nalp.utils import logging
@@ -18,16 +21,21 @@ class GumbelLSTMGenerator(LSTMGenerator):
     """
 
     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.0,
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (int): The amount of hidden neurons.
-            tau (float): Gumbel-Softmax temperature parameter.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            hidden_size: The amount of hidden neurons.
+            tau: Gumbel-Softmax temperature parameter.
 
         """
 
@@ -46,23 +54,23 @@ def __init__(
         logger.info("Class overrided.")
 
     @property
-    def tau(self):
-        """float: Gumbel-Softmax temperature parameter."""
+    def tau(self) -> float:
+        """Gumbel-Softmax temperature parameter."""
 
         return self._tau
 
     @tau.setter
-    def tau(self, tau):
+    def tau(self, tau: float) -> None:
         self._tau = tau
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            Logit-based predictions, Gumbel-Softmax outputs and predicted token.
+            (Tuple[tf.Tensor, tf.Tensor, tf.Tensor]): Logit-based predictions, Gumbel-Softmax outputs and predicted token.
 
         """
 
@@ -80,16 +88,18 @@ def call(self, x):
 
         return x, x_g, y_g
 
-    def generate_greedy_search(self, start, max_length=100):
+    def generate_greedy_search(
+        self, start: str, max_length: Optional[int] = 100
+    ) -> List[str]:
         """Generates text by using greedy search, where the sampled
         token is always sampled according to the maximum probability.
 
         Args:
-            start (str): The start string to generate the text.
-            max_length (int): Maximum length of generated text.
+            start: The start string to generate the text.
+            max_length: Maximum length of generated text.
 
         Returns:
-            A list holding the generated text.
+            (List[str]): Generated text.
 
         """
 
@@ -125,17 +135,22 @@ def generate_greedy_search(self, start, max_length=100):
 
         return sampled_tokens
 
-    def generate_temperature_sampling(self, start, max_length=100, temperature=1.0):
+    def generate_temperature_sampling(
+        self,
+        start: str,
+        max_length: Optional[int] = 100,
+        temperature: Optional[float] = 1.0,
+    ):
         """Generates text by using temperature sampling, where the sampled
         token is sampled according to a multinomial/categorical distribution.
 
         Args:
-            start (str): The start string to generate the text.
-            max_length (int): Length of generated text.
-            temperature (float): A temperature value to sample the token.
+            start: The start string to generate the text.
+            max_length: Length of generated text.
+            temperature: A temperature value to sample the token.
 
         Returns:
-            A list holding the generated text.
+            (List[str]): Generated text.
 
         """
 
@@ -177,19 +192,25 @@ def generate_temperature_sampling(self, start, max_length=100, temperature=1.0):
 
         return sampled_tokens
 
-    def generate_top_sampling(self, start, max_length=100, k=0, p=0.0):
+    def generate_top_sampling(
+        self,
+        start: str,
+        max_length: Optional[int] = 100,
+        k: Optional[int] = 0,
+        p: Optional[float] = 0.0,
+    ):
         """Generates text by using top-k and top-p sampling, where the sampled
         token is sampled according to the `k` most likely words distribution, as well
         as to the maximim cumulative probability `p`.
 
         Args:
-            start (str): The start string to generate the text.
-            max_length (int): Length of generated text.
-            k (int): Indicates the amount of likely words.
-            p (float): Maximum cumulative probability to be thresholded.
+            start: The start string to generate the text.
+            max_length: Length of generated text.
+            k: Indicates the amount of likely words.
+            p: Maximum cumulative probability to be thresholded.
 
         Returns:
-            A list holding the generated text.
+            (List[str]): Generated text.
 
         """
 
diff --git a/nalp/models/generators/gumbel_rmc.py b/nalp/models/generators/gumbel_rmc.py
index 095af81..230add0 100644
--- a/nalp/models/generators/gumbel_rmc.py
+++ b/nalp/models/generators/gumbel_rmc.py
@@ -1,9 +1,12 @@
 """Gumbel Relational Memory Core generator.
 """
 
+from typing import List, Optional, Tuple
+
 import tensorflow as tf
 
 import nalp.utils.constants as c
+from nalp.encoders.integer import IntegerEncoder
 from nalp.models.generators import RMCGenerator
 from nalp.models.layers import GumbelSoftmax
 from nalp.utils import logging
@@ -19,28 +22,28 @@ class GumbelRMCGenerator(RMCGenerator):
 
     def __init__(
         self,
-        encoder=None,
-        vocab_size=1,
-        embedding_size=32,
-        n_slots=3,
-        n_heads=5,
-        head_size=10,
-        n_blocks=1,
-        n_layers=3,
-        tau=5,
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: 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,
+        tau: Optional[float] = 5,
     ):
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            n_slots (int): Number of memory slots.
-            n_heads (int): Number of attention heads.
-            head_size (int): Size of each attention head.
-            n_blocks (int): Number of feed-forward networks.
-            n_layers (int): Amout of layers per feed-forward network.
-            tau (float): Gumbel-Softmax temperature parameter.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            n_slots: Number of memory slots.
+            n_heads: Number of attention heads.
+            head_size: Size of each attention head.
+            n_blocks: Number of feed-forward networks.
+            n_layers: Amout of layers per feed-forward network.
+            tau: Gumbel-Softmax temperature parameter.
 
         """
 
@@ -66,23 +69,23 @@ def __init__(
         logger.info("Class overrided.")
 
     @property
-    def tau(self):
-        """float: Gumbel-Softmax temperature parameter."""
+    def tau(self) -> float:
+        """Gumbel-Softmax temperature parameter."""
 
         return self._tau
 
     @tau.setter
-    def tau(self, tau):
+    def tau(self, tau: float) -> None:
         self._tau = tau
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            Logit-based predictions, Gumbel-Softmax outputs and predicted token.
+            (Tuple[tf.Tensor, tf.Tensor, tf.Tensor]): Logit-based predictions, Gumbel-Softmax outputs and predicted token.
 
         """
 
@@ -100,16 +103,18 @@ def call(self, x):
 
         return x, x_g, y_g
 
-    def generate_greedy_search(self, start, max_length=100):
+    def generate_greedy_search(
+        self, start: str, max_length: Optional[int] = 100
+    ) -> List[str]:
         """Generates text by using greedy search, where the sampled
         token is always sampled according to the maximum probability.
 
         Args:
-            start (str): The start string to generate the text.
-            max_length (int): Maximum length of generated text.
+            start: The start string to generate the text.
+            max_length: Maximum length of generated text.
 
         Returns:
-            A list holding the generated text.
+            (List[str]): Generated text.
 
         """
 
@@ -145,17 +150,22 @@ def generate_greedy_search(self, start, max_length=100):
 
         return sampled_tokens
 
-    def generate_temperature_sampling(self, start, max_length=100, temperature=1.0):
+    def generate_temperature_sampling(
+        self,
+        start: str,
+        max_length: Optional[int] = 100,
+        temperature: Optional[float] = 1.0,
+    ):
         """Generates text by using temperature sampling, where the sampled
         token is sampled according to a multinomial/categorical distribution.
 
         Args:
-            start (str): The start string to generate the text.
-            max_length (int): Length of generated text.
-            temperature (float): A temperature value to sample the token.
+            start: The start string to generate the text.
+            max_length: Length of generated text.
+            temperature: A temperature value to sample the token.
 
         Returns:
-            A list holding the generated text.
+            (List[str]): Generated text.
 
         """
 
@@ -197,19 +207,25 @@ def generate_temperature_sampling(self, start, max_length=100, temperature=1.0):
 
         return sampled_tokens
 
-    def generate_top_sampling(self, start, max_length=100, k=0, p=0.0):
+    def generate_top_sampling(
+        self,
+        start: str,
+        max_length: Optional[int] = 100,
+        k: Optional[int] = 0,
+        p: Optional[float] = 0.0,
+    ):
         """Generates text by using top-k and top-p sampling, where the sampled
         token is sampled according to the `k` most likely words distribution, as well
         as to the maximim cumulative probability `p`.
 
         Args:
-            start (str): The start string to generate the text.
-            max_length (int): Length of generated text.
-            k (int): Indicates the amount of likely words.
-            p (float): Maximum cumulative probability to be thresholded.
+            start: The start string to generate the text.
+            max_length: Length of generated text.
+            k: Indicates the amount of likely words.
+            p: Maximum cumulative probability to be thresholded.
 
         Returns:
-            A list holding the generated text.
+            (List[str]): Generated text.
 
         """
 
diff --git a/nalp/models/generators/linear.py b/nalp/models/generators/linear.py
index 8599412..cb1d997 100644
--- a/nalp/models/generators/linear.py
+++ b/nalp/models/generators/linear.py
@@ -1,6 +1,8 @@
 """Linear generator.
 """
 
+from typing import Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import Dense
 
@@ -16,14 +18,20 @@ class LinearGenerator(Generator):
 
     """
 
-    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 tensor.
-            noise_dim (int): Amount of noise dimensions.
-            n_samplings (int): Number of upsamplings to perform.
-            alpha (float): LeakyReLU activation threshold.
+            input_shape: An input shape for the tensor.
+            noise_dim: Amount of noise dimensions.
+            n_samplings: Number of upsamplings to perform.
+            alpha: LeakyReLU activation threshold.
 
         """
 
@@ -48,34 +56,34 @@ def __init__(self, input_shape=(784,), noise_dim=100, n_samplings=3, alpha=0.01)
         logger.info("Class overrided.")
 
     @property
-    def alpha(self):
-        """float: LeakyReLU activation threshold."""
+    def alpha(self) -> float:
+        """LeakyReLU activation threshold."""
 
         return self._alpha
 
     @alpha.setter
-    def alpha(self, alpha):
+    def alpha(self, alpha: float) -> None:
         self._alpha = alpha
 
     @property
-    def noise_dim(self):
-        """int: Amount of noise dimensions."""
+    def noise_dim(self) -> int:
+        """Amount of noise dimensions."""
 
         return self._noise_dim
 
     @noise_dim.setter
-    def noise_dim(self, noise_dim):
+    def noise_dim(self, noise_dim: int) -> None:
         self._noise_dim = noise_dim
 
-    def call(self, x, training=True):
+    def call(self, x: tf.Tensor, training: Optional[bool] = True) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            training (bool): Whether architecture is under training or not.
+            x: A tensorflow's tensor holding input data.
+            training: Whether architecture is under training or not.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/lstm.py b/nalp/models/generators/lstm.py
index dfbf05c..ef9b696 100644
--- a/nalp/models/generators/lstm.py
+++ b/nalp/models/generators/lstm.py
@@ -1,9 +1,13 @@
 """Long Short-Term Memory generator.
 """
 
+from typing import Optional
+
+import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, Embedding, LSTMCell
 
 from nalp.core import Generator
+from nalp.encoders.integer import IntegerEncoder
 from nalp.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -18,14 +22,20 @@ class LSTMGenerator(Generator):
 
     """
 
-    def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64):
+    def __init__(
+        self,
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        hidden_size: Optional[int] = 64,
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (int): The amount of hidden neurons.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            hidden_size: The amount of hidden neurons.
 
         """
 
@@ -53,23 +63,23 @@ def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64
         logger.info("Class overrided.")
 
     @property
-    def encoder(self):
-        """obj: An encoder generic object."""
+    def encoder(self) -> IntegerEncoder:
+        """An encoder generic object."""
 
         return self._encoder
 
     @encoder.setter
-    def encoder(self, encoder):
+    def encoder(self, encoder: IntegerEncoder) -> None:
         self._encoder = encoder
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/rmc.py b/nalp/models/generators/rmc.py
index b1794c3..b24a0ad 100644
--- a/nalp/models/generators/rmc.py
+++ b/nalp/models/generators/rmc.py
@@ -1,9 +1,12 @@
 """Relational Memory Core generator.
 """
+from typing import Optional
+
 import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, Embedding
 
 from nalp.core import Generator
+from nalp.encoders.integer import IntegerEncoder
 from nalp.models.layers.relational_memory_cell import RelationalMemoryCell
 from nalp.utils import logging
 
@@ -22,26 +25,26 @@ class RMCGenerator(Generator):
 
     def __init__(
         self,
-        encoder=None,
-        vocab_size=1,
-        embedding_size=32,
-        n_slots=3,
-        n_heads=5,
-        head_size=10,
-        n_blocks=1,
-        n_layers=3,
-    ):
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: 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,
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            n_slots (int): Number of memory slots.
-            n_heads (int): Number of attention heads.
-            head_size (int): Size of each attention head.
-            n_blocks (int): Number of feed-forward networks.
-            n_layers (int): Amout of layers per feed-forward network.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            n_slots: Number of memory slots.
+            n_heads: Number of attention heads.
+            head_size: Size of each attention head.
+            n_blocks: Number of feed-forward networks.
+            n_layers: Amout of layers per feed-forward network.
 
         """
 
@@ -71,23 +74,23 @@ def __init__(
         logger.info("Class overrided.")
 
     @property
-    def encoder(self):
-        """obj: An encoder generic object."""
+    def encoder(self) -> IntegerEncoder:
+        """An encoder generic object."""
 
         return self._encoder
 
     @encoder.setter
-    def encoder(self, encoder):
+    def encoder(self, encoder: IntegerEncoder) -> None:
         self._encoder = encoder
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/rnn.py b/nalp/models/generators/rnn.py
index e033a7c..79aaeb0 100644
--- a/nalp/models/generators/rnn.py
+++ b/nalp/models/generators/rnn.py
@@ -1,9 +1,13 @@
 """Recurrent Neural Network generator.
 """
 
+from typing import Optional
+
+import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, Embedding, SimpleRNNCell
 
 from nalp.core import Generator
+from nalp.encoders.integer import IntegerEncoder
 from nalp.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -18,14 +22,20 @@ class RNNGenerator(Generator):
 
     """
 
-    def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64):
+    def __init__(
+        self,
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        hidden_size: Optional[int] = 64,
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (int): The amount of hidden neurons.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            hidden_size: The amount of hidden neurons.
 
         """
 
@@ -53,23 +63,23 @@ def __init__(self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=64
         logger.info("Class overrided.")
 
     @property
-    def encoder(self):
-        """obj: An encoder generic object."""
+    def encoder(self) -> IntegerEncoder:
+        """An encoder generic object."""
 
         return self._encoder
 
     @encoder.setter
-    def encoder(self, encoder):
+    def encoder(self, encoder: IntegerEncoder) -> None:
         self._encoder = encoder
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/generators/stacked_rnn.py b/nalp/models/generators/stacked_rnn.py
index 0116c49..0ff171f 100644
--- a/nalp/models/generators/stacked_rnn.py
+++ b/nalp/models/generators/stacked_rnn.py
@@ -1,9 +1,13 @@
 """Stacked Recurrent Neural Network generator.
 """
 
+from typing import Optional, Tuple
+
+import tensorflow as tf
 from tensorflow.keras.layers import RNN, Dense, Embedding, SimpleRNNCell
 
 from nalp.core import Generator
+from nalp.encoders.integer import IntegerEncoder
 from nalp.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -19,15 +23,19 @@ class StackedRNNGenerator(Generator):
     """
 
     def __init__(
-        self, encoder=None, vocab_size=1, embedding_size=32, hidden_size=(64, 64)
-    ):
+        self,
+        encoder: Optional[IntegerEncoder] = None,
+        vocab_size: Optional[int] = 1,
+        embedding_size: Optional[int] = 32,
+        hidden_size: Optional[Tuple[int, ...]] = (64, 64),
+    ) -> None:
         """Initialization method.
 
         Args:
-            encoder (IntegerEncoder): An index to vocabulary encoder.
-            vocab_size (int): The size of the vocabulary.
-            embedding_size (int): The size of the embedding layer.
-            hidden_size (tuple): Amount of hidden neurons per cell.
+            encoder: An index to vocabulary encoder.
+            vocab_size: The size of the vocabulary.
+            embedding_size: The size of the embedding layer.
+            hidden_size: Amount of hidden neurons per cell.
 
         """
 
@@ -59,23 +67,23 @@ def __init__(
         logger.info("Class overrided.")
 
     @property
-    def encoder(self):
-        """obj: An encoder generic object."""
+    def encoder(self) -> IntegerEncoder:
+        """An encoder generic object."""
 
         return self._encoder
 
     @encoder.setter
-    def encoder(self, encoder):
+    def encoder(self, encoder: IntegerEncoder) -> None:
         self._encoder = encoder
 
-    def call(self, x):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
+            x: A tensorflow's tensor holding input data.
 
         Returns:
-            The same tensor after passing through each defined layer.
+            (tf.Tensor): The same tensor after passing through each defined layer.
 
         """
 
diff --git a/nalp/models/layers/gumbel_softmax.py b/nalp/models/layers/gumbel_softmax.py
index 2c1410d..08f83ef 100644
--- a/nalp/models/layers/gumbel_softmax.py
+++ b/nalp/models/layers/gumbel_softmax.py
@@ -1,20 +1,22 @@
 """Gumbel-Softmax layer.
 """
 
+from typing import Any, Dict, Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import Layer
 
 import nalp.utils.constants as c
 
 
-def gumbel_distribution(input_shape):
+def gumbel_distribution(input_shape: Tuple[int, ...]) -> tf.Tensor:
     """Samples a tensor from a Gumbel distribution.
 
     Args:
-        input_shape (tuple): Shape of tensor to be sampled.
+        input_shape: Shape of tensor to be sampled.
 
     Returns:
-        An input_shape tensor sampled from a Gumbel distribution.
+        (tf.Tensor): An input_shape tensor sampled from a Gumbel distribution.
 
     """
 
@@ -38,11 +40,11 @@ class GumbelSoftmax(Layer):
 
     """
 
-    def __init__(self, axis=-1, **kwargs):
+    def __init__(self, axis: Optional[int] = -1, **kwargs) -> None:
         """Initialization method.
 
         Args:
-            axis (int): Axis to perform the softmax operation.
+            axis: Axis to perform the softmax operation.
 
         """
 
@@ -51,15 +53,15 @@ def __init__(self, axis=-1, **kwargs):
         # Defining a property for holding the intended axis
         self.axis = axis
 
-    def call(self, inputs, tau):
+    def call(self, inputs: tf.Tensor, tau: float) -> Tuple[tf.Tensor, tf.Tensor]:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            x (tf.tensor): A tensorflow's tensor holding input data.
-            tau (float): Gumbel-Softmax temperature parameter.
+            x: A tensorflow's tensor holding input data.
+            tau: Gumbel-Softmax temperature parameter.
 
         Returns:
-            Gumbel-Softmax output and its argmax token.
+            (Tuple[tf.Tensor, tf.Tensor]): Gumbel-Softmax output and its argmax token.
 
         """
 
@@ -74,8 +76,13 @@ def call(self, inputs, tau):
 
         return x, y
 
-    def get_config(self):
-        """Gets the configuration of the layer for further serialization."""
+    def get_config(self) -> Dict[str, Any]:
+        """Gets the configuration of the layer for further serialization.
+
+        Returns:
+            (Dict[str, Any]): Configuration dictionary.
+
+        """
 
         config = {"axis": self.axis}
         base_config = super(GumbelSoftmax, self).get_config()
diff --git a/nalp/models/layers/multi_head_attention.py b/nalp/models/layers/multi_head_attention.py
index 5d204d8..21fce8a 100644
--- a/nalp/models/layers/multi_head_attention.py
+++ b/nalp/models/layers/multi_head_attention.py
@@ -1,24 +1,28 @@
 """Multi-Head Attention layer.
 """
 
+from typing import Any, Dict, Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import Dense, Layer
 
 import nalp.utils.constants as c
 
 
-def scaled_dot_product_attention(q, k, v, mask):
+def scaled_dot_product_attention(
+    q: tf.Tensor, k: tf.Tensor, v: tf.Tensor, mask: tf.Tensor
+) -> Tuple[tf.Tensor, tf.Tensor]:
     """Calculate the attention weights, such that q, k, v must have matching
     leading dimensions, and k, v must have matching penultimate dimension.
 
     Args:
-        q (tf.tensor): Query tensor.
-        k (tf.tensor): Key tensor.
-        v (tf.tensor): Value tensor.
-        mask (tf.tensor): Mask to be applied.
+        q: Query tensor.
+        k: Key tensor.
+        v: Value tensor.
+        mask: Mask to be applied.
 
     Returns:
-        An attention-based output tensor and its attention weights.
+        (Tuple[tf.Tensor, tf.Tensor]): An attention-based output tensor and its attention weights.
 
     """
 
@@ -53,12 +57,12 @@ class MultiHeadAttention(Layer):
 
     """
 
-    def __init__(self, n_features, n_heads, **kwargs):
+    def __init__(self, n_features: int, n_heads: int, **kwargs) -> None:
         """Initialization method.
 
         Args:
-            n_features (int): Number of input features.
-            n_heads (int): Number of attention heads.
+            n_features: Number of input features.
+            n_heads: Number of attention heads.
 
         """
 
@@ -88,14 +92,14 @@ def __init__(self, n_features, n_heads, **kwargs):
         # Creating the final linear layer
         self.out = Dense(n_features)
 
-    def _split_heads(self, x):
+    def _split_heads(self, x: tf.Tensor) -> tf.Tensor:
         """Split the last tensor dimension into (n_heads, depth) and transposes its result.
 
         Args:
-            x (tf.tensor): Tensor to be splitted and transposed.
+            x: Tensor to be splitted and transposed.
 
         Returns:
-            Splitted and transposed tensor into shape equal to (batch_size, n_heads, length, depth).
+            (tf.Tensor): Splitted and transposed tensor into shape equal to (batch_size, n_heads, length, depth).
 
         """
 
@@ -104,17 +108,19 @@ def _split_heads(self, x):
 
         return tf.transpose(x, perm=[0, 2, 1, 3])
 
-    def call(self, q, k, v, mask=None):
+    def call(
+        self, q: tf.Tensor, k: tf.Tensor, v: tf.Tensor, mask: Optional[tf.Tensor] = None
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            q (tf.tensor): Query tensor.
-            k (tf.tensor): Key tensor.
-            v (tf.tensor): Value tensor.
-            mask (tf.tensor): Mask to be applied.
+            q: Query tensor.
+            k: Key tensor.
+            v: Value tensor.
+            mask: Mask to be applied.
 
         Returns:
-            An attention-based output tensor and its attention weights.
+            (Tuple[tf.Tensor, tf.Tensor]): An attention-based output tensor and its attention weights.
 
         """
 
@@ -142,8 +148,13 @@ def call(self, q, k, v, mask=None):
 
         return output, attn_weights
 
-    def get_config(self):
-        """Gets the configuration of the layer for further serialization."""
+    def get_config(self) -> Dict[str, Any]:
+        """Gets the configuration of the layer for further serialization.
+
+        Returns:
+            (Dict[str, Any]): Configuration dictionary.
+
+        """
 
         config = {
             "n_features": self.n_features,
diff --git a/nalp/models/layers/relational_memory_cell.py b/nalp/models/layers/relational_memory_cell.py
index b09b13a..90d2e61 100644
--- a/nalp/models/layers/relational_memory_cell.py
+++ b/nalp/models/layers/relational_memory_cell.py
@@ -1,6 +1,8 @@
 """Relational-Memory Cell layer.
 """
 
+from typing import Any, Dict, List, Optional, Tuple
+
 import tensorflow as tf
 from tensorflow.keras.layers import (
     AbstractRNNCell,
@@ -22,45 +24,45 @@ class RelationalMemoryCell(AbstractRNNCell):
 
     def __init__(
         self,
-        n_slots,
-        n_heads,
-        head_size,
-        n_blocks=1,
-        n_layers=3,
-        activation="tanh",
-        recurrent_activation="hard_sigmoid",
-        forget_bias=1.0,
-        kernel_initializer="glorot_uniform",
-        recurrent_initializer="orthogonal",
-        bias_initializer="zeros",
-        kernel_regularizer=None,
-        recurrent_regularizer=None,
-        bias_regularizer=None,
-        kernel_constraint=None,
-        recurrent_constraint=None,
-        bias_constraint=None,
+        n_slots: int,
+        n_heads: int,
+        head_size: int,
+        n_blocks: Optional[int] = 1,
+        n_layers: Optional[int] = 3,
+        activation: Optional[str] = "tanh",
+        recurrent_activation: Optional[str] = "hard_sigmoid",
+        forget_bias: Optional[float] = 1.0,
+        kernel_initializer: Optional[str] = "glorot_uniform",
+        recurrent_initializer: Optional[str] = "orthogonal",
+        bias_initializer: Optional[str] = "zeros",
+        kernel_regularizer: Optional[str] = None,
+        recurrent_regularizer: Optional[str] = None,
+        bias_regularizer: Optional[str] = None,
+        kernel_constraint: Optional[str] = None,
+        recurrent_constraint: Optional[str] = None,
+        bias_constraint: Optional[str] = None,
         **kwargs
     ):
         """Initialization method.
 
         Args:
-            n_slots (int): Number of memory slots.
-            n_heads (int): Number of attention heads.
-            head_size (int): Size of each attention head.
-            n_blocks (int): Number of feed-forward networks.
-            n_layers (int): Amout of layers per feed-forward network.
-            activation (str): Output activation function.
-            recurrent_activation (str): Recurrent step activation function.
-            forget_bias (float): Forget gate bias values.
-            kernel_initializer (str): Kernel initializer function.
-            recurrent_initializer (str): Recurrent kernel initializer function.
-            bias_initializer (str): Bias initializer function.
-            kernel_regularizer (str): Kernel regularizer function.
-            recurrent_regularizer (str): Recurrent kernel regularizer function.
-            bias_regularizer (str): Bias regularizer function.
-            kernel_constraint (str): Kernel constraint function.
-            recurrent_constraint (str): Recurrent kernel constraint function.
-            bias_constraint (str): Bias constraint function.
+            n_slots: Number of memory slots.
+            n_heads: Number of attention heads.
+            head_size: Size of each attention head.
+            n_blocks: Number of feed-forward networks.
+            n_layers: Amout of layers per feed-forward network.
+            activation: Output activation function.
+            recurrent_activation: Recurrent step activation function.
+            forget_bias: Forget gate bias values.
+            kernel_initializer: Kernel initializer function.
+            recurrent_initializer: Recurrent kernel initializer function.
+            bias_initializer: Bias initializer function.
+            kernel_regularizer: Kernel regularizer function.
+            recurrent_regularizer: Recurrent kernel regularizer function.
+            bias_regularizer: Bias regularizer function.
+            kernel_constraint: Kernel constraint function.
+            recurrent_constraint: Recurrent kernel constraint function.
+            bias_constraint: Bias constraint function.
 
         """
 
@@ -121,18 +123,18 @@ def __init__(
         self.attn = MultiHeadAttention(self.slot_size, self.n_heads)
 
     @property
-    def state_size(self):
+    def state_size(self) -> List[int, int]:
         return [self.units, self.units]
 
     @property
-    def output_size(self):
+    def output_size(self) -> int:
         return self.units
 
-    def build(self, input_shape):
+    def build(self, input_shape: tf.Tensor) -> None:
         """Builds up the cell according to its input shape.
 
         Args:
-            input_shape (tf.tensor): Tensor holding the input shape.
+            input_shape: Tensor holding the input shape.
 
         """
 
@@ -166,15 +168,15 @@ def build(self, input_shape):
         # Marking the built property as `True`
         self.built = True
 
-    def _attend_over_memory(self, inputs, memory):
+    def _attend_over_memory(self, inputs: tf.Tensor, memory: tf.Tensor) -> tf.Tensor:
         """Performs an Attention mechanism over the current memory.
 
         Args:
-            inputs (tf.tensor): An input tensor.
-            memory (tf.tensor): Current memory tensor.
+            inputs: An input tensor.
+            memory: Current memory tensor.
 
         Returns:
-            Updated current memory based on Multi-Head Attention mechanism.
+            (tf.Tensor): Updated current memory based on Multi-Head Attention mechanism.
 
         """
 
@@ -204,15 +206,17 @@ def _attend_over_memory(self, inputs, memory):
 
         return memory
 
-    def call(self, inputs, states):
+    def call(
+        self, inputs: tf.Tensor, states: List[tf.Tensor]
+    ) -> Tuple[tf.Tensor, List[tf.Tensor]]:
         """Method that holds vital information whenever this class is called.
 
         Args:
-            inputs (tf.tensor): An input tensor.
-            states (list): A list holding previous states and memories.
+            inputs: An input tensor.
+            states: A list holding previous states and memories.
 
         Returns:
-            Output states as well as current state and memory.
+            (Tuple[tf.Tensor, List[tf.Tensor]]): Output states as well as current state and memory.
 
         """
 
@@ -270,16 +274,21 @@ def call(self, inputs, states):
 
         return h, [h, m]
 
-    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
+    def get_initial_state(
+        self,
+        inputs: Optional[tf.Tensor] = None,
+        batch_size: Optional[int] = None,
+        dtype: Optional[tf.DType] = None,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
         """Gets the cell initial state by creating an identity matrix.
 
         Args:
-            inputs (tf.tensor): An input tensor.
-            batch_size (int): Size of the batch.
-            dtype (dtype): Dtype from input tensor.
+            inputs: An input tensor.
+            batch_size: Size of the batch.
+            dtype: Dtype from input tensor.
 
         Returns:
-            Initial states.
+            (Tuple[tf.Tensor, tf.Tensor]): Initial states.
 
         """
 
@@ -307,8 +316,13 @@ def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
 
         return states, states
 
-    def get_config(self):
-        """Gets the configuration of the layer for further serialization."""
+    def get_config(self) -> Dict[str, Any]:
+        """Gets the configuration of the layer for further serialization.
+
+        Returns:
+            (Dict[str, Any]): Configuration dictionary.
+
+        """
 
         config = {
             "n_slots": self.n_slots,