text: """
  Implementation of the RNN model
  """
  from typing import List, Tuple
  import numpy as np
  import torch
  import torch.nn as tnn
  import torch.nn.functional as tnnf
  
  from reinvent_models.reinvent_core.models import vocabulary as mv
  
  # from models import vocabulary as mv
  
  
  class RNN(tnn.Module):
      """
      Implements a N layer GRU(M) cell including an embedding layer
      and an output linear layer back to the size of the vocabulary
      """
  
      def __init__(
          self,
          voc_size,
          layer_size=512,
          num_layers=3,
          cell_type="gru",
          embedding_layer_size=256,
          dropout=0.0,
          layer_normalization=False,
      ):
          """
          Implements a N layer GRU|LSTM cell including an embedding layer and an output linear layer back to the size of the
          vocabulary
          :param voc_size: Size of the vocabulary.
          :param layer_size: Size of each of the RNN layers.
          :param num_layers: Number of RNN layers.
          :param embedding_layer_size: Size of the embedding layer.
          """
          super(RNN, self).__init__()
  
          self._layer_size = layer_size
          self._embedding_layer_size = embedding_layer_size
          self._num_layers = num_layers
          self._cell_type = cell_type.lower()
          self._dropout = dropout
          self._layer_normalization = layer_normalization
  
          self._embedding = tnn.Embedding(voc_size, self._embedding_layer_size)
          if self._cell_type == "gru":
              self._rnn = tnn.GRU(
                  self._embedding_layer_size,
                  self._layer_size,
                  num_layers=self._num_layers,
                  dropout=self._dropout,
                  batch_first=True,
              )
          elif self._cell_type == "lstm":
              self._rnn = tnn.LSTM(
                  self._embedding_layer_size,
                  self._layer_size,
                  num_layers=self._num_layers,
                  dropout=self._dropout,
                  batch_first=True,
              )
          else:
              raise ValueError(
                  'Value of the parameter cell_type should be "gru" or "lstm"'
              )
          self._linear = tnn.Linear(self._layer_size, voc_size)
  
      def forward(self, input_vector, hidden_state=None):  # pylint: disable=W0221
          """
          Performs a forward pass on the model. Note: you pass the **whole** sequence.
          :param input_vector: Input tensor (batch_size, seq_size).
          :param hidden_state: Hidden state tensor.
          """
          batch_size, seq_size = input_vector.size()
          if hidden_state is None:
              size = (self._num_layers, batch_size, self._layer_size)
              if self._cell_type == "gru":
                  hidden_state = torch.zeros(*size)
              else:
                  hidden_state = [torch.zeros(*size), torch.zeros(*size)]
          embedded_data = self._embedding(input_vector)  # (batch,seq,embedding)
          output_vector, hidden_state_out = self._rnn(embedded_data, hidden_state)
  
          if self._layer_normalization:
              output_vector = tnnf.layer_norm(output_vector, output_vector.size()[1:])
          output_vector = output_vector.reshape(-1, self._layer_size)
  
          output_data = self._linear(output_vector).view(batch_size, seq_size, -1)
  
          return output_data, hidden_state_out
  
      def get_params(self):
          """
          Returns the configuration parameters of the model.
          """
          return {
              "dropout": self._dropout,
              "layer_size": self._layer_size,
              "num_layers": self._num_layers,
              "cell_type": self._cell_type,
              "embedding_layer_size": self._embedding_layer_size,
          }
  
  
  class Model:
      """
      Implements an RNN model using SMILES.
      """
  
      def __init__(
          self,
          vocabulary: mv.Vocabulary,
          tokenizer,
          network_params=None,
          max_sequence_length=256,
          no_cuda=False,
      ):
          """
          Implements an RNN.
          :param vocabulary: Vocabulary to use.
          :param tokenizer: Tokenizer to use.
          :param network_params: Dictionary with all parameters required to correctly initialize the RNN class.
          :param max_sequence_length: The max size of SMILES sequence that can be generated.
          """
          self.vocabulary = vocabulary
          self.tokenizer = tokenizer
          self.max_sequence_length = max_sequence_length
  
          if not isinstance(network_params, dict):
              network_params = {}
  
          self.network = RNN(len(self.vocabulary), **network_params)
          if torch.cuda.is_available() and not no_cuda:
              self.network.cuda()
  
          self._nll_loss = tnn.NLLLoss(reduction="none")
  
      @classmethod
      def load_from_file(cls, file_path: str, sampling_mode=False):
          """
          Loads a model from a single file
          :param file_path: input file path
          :return: new instance of the RNN or an exception if it was not possible to load it.
          """
          if torch.cuda.is_available():
              save_dict = torch.load(file_path)
          else:
              save_dict = torch.load(file_path, map_location=lambda storage, loc: storage)
  
          network_params = save_dict.get("network_params", {})
          model = Model(
              vocabulary=save_dict["vocabulary"],
              tokenizer=save_dict.get("tokenizer", mv.SMILESTokenizer()),
              network_params=network_params,
              max_sequence_length=save_dict["max_sequence_length"],
          )
          model.network.load_state_dict(save_dict["network"])
          if sampling_mode:
              model.network.eval()
          return model
  
      def save(self, file: str):
          """
          Saves the model into a file
          :param file: it's actually a path
          """
          save_dict = {
              "vocabulary": self.vocabulary,
              "tokenizer": self.tokenizer,
              "max_sequence_length": self.max_sequence_length,
              "network": self.network.state_dict(),
              "network_params": self.network.get_params(),
          }
          torch.save(save_dict, file)
  
      def likelihood_smiles(self, smiles) -> torch.Tensor:
          tokens = [self.tokenizer.tokenize(smile) for smile in smiles]
          encoded = [self.vocabulary.encode(token) for token in tokens]
          sequences = [torch.tensor(encode, dtype=torch.long) for encode in encoded]
  
          def collate_fn(encoded_seqs):
              """Function to take a list of encoded sequences and turn them into a batch"""
              max_length = max([seq.size(0) for seq in encoded_seqs])
              collated_arr = torch.zeros(
                  len(encoded_seqs), max_length, dtype=torch.long
              )  # padded with zeroes
              for i, seq in enumerate(encoded_seqs):
                  collated_arr[i, : seq.size(0)] = seq
              return collated_arr
  
          padded_sequences = collate_fn(sequences)
          return self.likelihood(padded_sequences)
  
      def likelihood(self, sequences) -> torch.Tensor:
          """
          Retrieves the likelihood of a given sequence. Used in training.
  
          :param sequences: (batch_size, sequence_length) A batch of sequences
          :return:  (batch_size) Log likelihood for each example.
          """
          logits, _ = self.network(sequences[:, :-1])  # all steps done at once
          log_probs = logits.log_softmax(dim=2)
          return self._nll_loss(log_probs.transpose(1, 2), sequences[:, 1:]).sum(dim=1)
  
      def sample_smiles(self, num=128, batch_size=128) -> Tuple[List, np.array]:
          """
          Samples n SMILES from the model.
          :param num: Number of SMILES to sample.
          :param batch_size: Number of sequences to sample at the same time.
          :return:
              :smiles: (n) A list with SMILES.
              :likelihoods: (n) A list of likelihoods.
          """
          batch_sizes = [batch_size for _ in range(num // batch_size)] + [
              num % batch_size
          ]
          smiles_sampled = []
          likelihoods_sampled = []
  
          for size in batch_sizes:
              if not size:
                  break
              seqs, likelihoods = self._sample(batch_size=size)
              smiles = [
                  self.tokenizer.untokenize(self.vocabulary.decode(seq))
                  for seq in seqs.cpu().numpy()
              ]
  
              smiles_sampled.extend(smiles)
              likelihoods_sampled.append(likelihoods.data.cpu().numpy())
  
              del seqs, likelihoods
          return smiles_sampled, np.concatenate(likelihoods_sampled)
  
      def sample_sequences_and_smiles(
          self, batch_size=128
      ) -> Tuple[torch.Tensor, List, torch.Tensor]:
          seqs, likelihoods = self._sample(batch_size=batch_size)
          smiles = [
              self.tokenizer.untokenize(self.vocabulary.decode(seq))
              for seq in seqs.cpu().numpy()
          ]
          return seqs, smiles, likelihoods
  
      # @torch.no_grad()
      def _sample(self, batch_size=128) -> Tuple[torch.Tensor, torch.Tensor]:
          start_token = torch.zeros(batch_size, dtype=torch.long)
          start_token[:] = self.vocabulary["^"]
          input_vector = start_token
          sequences = [
              self.vocabulary["^"] * torch.ones([batch_size, 1], dtype=torch.long)
          ]
          # NOTE: The first token never gets added in the loop so the sequences are initialized with a start token
          hidden_state = None
          nlls = torch.zeros(batch_size)
          for _ in range(self.max_sequence_length - 1):
              logits, hidden_state = self.network(input_vector.unsqueeze(1), hidden_state)
              logits = logits.squeeze(1)
              probabilities = logits.softmax(dim=1)
              log_probs = logits.log_softmax(dim=1)
              input_vector = torch.multinomial(probabilities, 1).view(-1)
              sequences.append(input_vector.view(-1, 1))
              nlls += self._nll_loss(log_probs, input_vector)
              if input_vector.sum() == 0:
                  break
  
          sequences = torch.cat(sequences, 1)
          return sequences.data, nlls
  
      def get_network_parameters(self):
          return self.network.parameters()
  