<commit_before>from typing import Tuple, List, Union, Callable, NamedTuple

import tensorflow as tf
from typeguard import check_argument_types

from neuralmonkey.model.stateful import (
    TemporalStatefulWithOutput, TemporalStateful)
from neuralmonkey.model.parameterized import InitializerSpecs
from neuralmonkey.model.model_part import ModelPart
from neuralmonkey.nn.ortho_gru_cell import OrthoGRUCell, NematusGRUCell
from neuralmonkey.nn.utils import dropout
from neuralmonkey.vocabulary import Vocabulary
from neuralmonkey.decorators import tensor
from neuralmonkey.model.sequence import (
    EmbeddedSequence, EmbeddedFactorSequence)
from neuralmonkey.tf_utils import layer_norm

RNN_CELL_TYPES = {
    "NematusGRU": NematusGRUCell,
    "GRU": OrthoGRUCell,
    "LSTM": tf.nn.rnn_cell.LSTMCell
}

RNN_DIRECTIONS = ["forward", "backward", "bidirectional"]

# pylint: disable=invalid-name
RNNSpecTuple = Union[Tuple[int], Tuple[int, str], Tuple[int, str, str]]
RNNCellTuple = Tuple[tf.nn.rnn_cell.RNNCell, tf.nn.rnn_cell.RNNCell]
# pylint: enable=invalid-name


class RNNSpec(NamedTuple(
        "RNNSpec",
        [("size", int),
         ("direction", str),
         ("cell_type", str)])):
    """Recurrent neural network specifications.

    Attributes:
        size: The state size.
        direction: The RNN processing direction. One of ``forward``,
            ``backward``, and ``bidirectional``.
        cell_type: The recurrent cell type to use. Refer to
            ``encoders.recurrent.RNN_CELL_TYPES`` for possible values.
    """


def _make_rnn_spec(size: int,
                   direction: str = "bidirectional",
                   cell_type: str = "GRU") -> RNNSpec:
    if size <= 0:
        raise ValueError(
            "RNN size must be a positive integer. {} given.".format(size))

    if direction not in RNN_DIRECTIONS:
        raise ValueError("RNN direction must be one of {}. {} given."
                         .format(str(RNN_DIRECTIONS), direction))

    if cell_type not in RNN_CELL_TYPES:
        raise ValueError("RNN cell type must be one of {}. {} given."
                         .format(str(RNN_CELL_TYPES), cell_type))

    return RNNSpec(size, direction, cell_type)


def _make_rnn_cell(spec: RNNSpec) -> Callable[[], tf.nn.rnn_cell.RNNCell]:
    """Return the graph template for creating RNN cells."""
    return RNN_CELL_TYPES[spec.cell_type](spec.size)


def rnn_layer(rnn_input: tf.Tensor,
              lengths: tf.Tensor,
              rnn_spec: RNNSpec) -> Tuple[tf.Tensor, tf.Tensor]:
    """Construct a RNN layer given its inputs and specs.

    Arguments:
        rnn_inputs: The input sequence to the RNN.
        lengths: Lengths of input sequences.
        rnn_spec: A valid RNNSpec tuple specifying the network architecture.
        add_residual: Add residual connections to the layer output.
    """
    if rnn_spec.direction == "bidirectional":
        fw_cell = _make_rnn_cell(rnn_spec)
        bw_cell = _make_rnn_cell(rnn_spec)

        outputs_tup, states_tup = tf.nn.bidirectional_dynamic_rnn(
            fw_cell, bw_cell, rnn_input, sequence_length=lengths,
            dtype=tf.float32)

        outputs = tf.concat(outputs_tup, 2)

        if rnn_spec.cell_type == "LSTM":
            states_tup = (state.h for state in states_tup)

        final_state = tf.concat(list(states_tup), 1)
    else:
        if rnn_spec.direction == "backward":
            rnn_input = tf.reverse_sequence(rnn_input, lengths, seq_axis=1)

        cell = _make_rnn_cell(rnn_spec)
        outputs, final_state = tf.nn.dynamic_rnn(
            cell, rnn_input, sequence_length=lengths, dtype=tf.float32)

        if rnn_spec.direction == "backward":
            outputs = tf.reverse_sequence(outputs, lengths, seq_axis=1)

        if rnn_spec.cell_type == "LSTM":
            final_state = final_state.h

    return outputs, final_state


class RecurrentEncoder(ModelPart, TemporalStatefulWithOutput):

    # pylint: disable=too-many-arguments
    def __init__(self,
                 name: str,
                 input_sequence: TemporalStateful,
                 rnn_layers: List[RNNSpecTuple],
                 add_residual: bool = False,
                 add_layer_norm: bool = False,
                 dropout_keep_prob: float = 1.0,
                 reuse: ModelPart = None,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None,
                 initializers: InitializerSpecs = None) -> None:
        """Create a new instance of a recurrent encoder.

        Arguments:
            name: ModelPart name.
            input_seqeunce: The input sequence for the encoder.
            rnn_size: The dimension of the RNN hidden state vector.
            rnn_cell: One of "GRU", "NematusGRU", "LSTM". Which kind of memory
                cell to use.
            rnn_direction: One of "forward", "backward", "bidirectional". In
                what order to process the input sequence. Note that choosing
                "bidirectional" will double the resulting vector dimension as
                well as the number of encoder parameters.
            add_residual: Add residual connections to the RNN layer output.
            add_layer_norm: Add layer normalization after each RNN layer.
            dropout_keep_prob: 1 - dropout probability.
            save_checkpoint: ModelPart save checkpoint file.
            load_checkpoint: ModelPart load checkpoint file.
        """
        check_argument_types()
        ModelPart.__init__(self, name, reuse, save_checkpoint, load_checkpoint,
                           initializers)
        TemporalStatefulWithOutput.__init__(self)

        self.input_sequence = input_sequence
        self.dropout_keep_prob = dropout_keep_prob
        self.rnn_specs = [_make_rnn_spec(*r) for r in rnn_layers]
        self.add_residual = add_residual
        self.add_layer_norm = add_layer_norm

        if self.dropout_keep_prob <= 0.0 or self.dropout_keep_prob > 1.0:
            raise ValueError("Dropout keep prob must be inside (0,1].")

        self._variable_scope.set_initializer(
            tf.random_normal_initializer(stddev=0.001))
    # pylint: enable=too-many-arguments

    @tensor
    def rnn_input(self) -> tf.Tensor:
        return dropout(self.input_sequence.temporal_states,
                       self.dropout_keep_prob, self.train_mode)

    @tensor
    def rnn(self) -> Tuple[tf.Tensor, tf.Tensor]:
        layer_input = self.rnn_input  # type: tf.Tensor
        layer_final = self.rnn_input[:, -1]

        for i, rnn_spec in enumerate(self.rnn_specs):
            with tf.variable_scope("rnn_{}_{}".format(i, rnn_spec.direction),
                                   reuse=tf.AUTO_REUSE):

                if self.add_layer_norm:
                    layer_input = layer_norm(layer_input)

                layer_output, layer_final_output = rnn_layer(
                    layer_input, self.input_sequence.lengths, rnn_spec)

                layer_output = dropout(
                    layer_output, self.dropout_keep_prob, self.train_mode)
                layer_final_output = dropout(
                    layer_final_output, self.dropout_keep_prob,
                    self.train_mode)

                in_dim = layer_input.get_shape()[-1]<commit_msg>pylint fix<commit_after>from typing import Tuple, List, Union, Callable, NamedTuple

import tensorflow as tf
from typeguard import check_argument_types

from neuralmonkey.model.stateful import (
    TemporalStatefulWithOutput, TemporalStateful)
from neuralmonkey.model.parameterized import InitializerSpecs
from neuralmonkey.model.model_part import ModelPart
from neuralmonkey.nn.ortho_gru_cell import OrthoGRUCell, NematusGRUCell
from neuralmonkey.nn.utils import dropout
from neuralmonkey.vocabulary import Vocabulary
from neuralmonkey.decorators import tensor
from neuralmonkey.model.sequence import (
    EmbeddedSequence, EmbeddedFactorSequence)
from neuralmonkey.tf_utils import layer_norm

RNN_CELL_TYPES = {
    "NematusGRU": NematusGRUCell,
    "GRU": OrthoGRUCell,
    "LSTM": tf.nn.rnn_cell.LSTMCell
}

RNN_DIRECTIONS = ["forward", "backward", "bidirectional"]

# pylint: disable=invalid-name
RNNSpecTuple = Union[Tuple[int], Tuple[int, str], Tuple[int, str, str]]
RNNCellTuple = Tuple[tf.nn.rnn_cell.RNNCell, tf.nn.rnn_cell.RNNCell]
# pylint: enable=invalid-name


class RNNSpec(NamedTuple(
        "RNNSpec",
        [("size", int),
         ("direction", str),
         ("cell_type", str)])):
    """Recurrent neural network specifications.

    Attributes:
        size: The state size.
        direction: The RNN processing direction. One of ``forward``,
            ``backward``, and ``bidirectional``.
        cell_type: The recurrent cell type to use. Refer to
            ``encoders.recurrent.RNN_CELL_TYPES`` for possible values.
    """


def _make_rnn_spec(size: int,
                   direction: str = "bidirectional",
                   cell_type: str = "GRU") -> RNNSpec:
    if size <= 0:
        raise ValueError(
            "RNN size must be a positive integer. {} given.".format(size))

    if direction not in RNN_DIRECTIONS:
        raise ValueError("RNN direction must be one of {}. {} given."
                         .format(str(RNN_DIRECTIONS), direction))

    if cell_type not in RNN_CELL_TYPES:
        raise ValueError("RNN cell type must be one of {}. {} given."
                         .format(str(RNN_CELL_TYPES), cell_type))

    return RNNSpec(size, direction, cell_type)


def _make_rnn_cell(spec: RNNSpec) -> Callable[[], tf.nn.rnn_cell.RNNCell]:
    """Return the graph template for creating RNN cells."""
    return RNN_CELL_TYPES[spec.cell_type](spec.size)


def rnn_layer(rnn_input: tf.Tensor,
              lengths: tf.Tensor,
              rnn_spec: RNNSpec) -> Tuple[tf.Tensor, tf.Tensor]:
    """Construct a RNN layer given its inputs and specs.

    Arguments:
        rnn_inputs: The input sequence to the RNN.
        lengths: Lengths of input sequences.
        rnn_spec: A valid RNNSpec tuple specifying the network architecture.
        add_residual: Add residual connections to the layer output.
    """
    if rnn_spec.direction == "bidirectional":
        fw_cell = _make_rnn_cell(rnn_spec)
        bw_cell = _make_rnn_cell(rnn_spec)

        outputs_tup, states_tup = tf.nn.bidirectional_dynamic_rnn(
            fw_cell, bw_cell, rnn_input, sequence_length=lengths,
            dtype=tf.float32)

        outputs = tf.concat(outputs_tup, 2)

        if rnn_spec.cell_type == "LSTM":
            states_tup = (state.h for state in states_tup)

        final_state = tf.concat(list(states_tup), 1)
    else:
        if rnn_spec.direction == "backward":
            rnn_input = tf.reverse_sequence(rnn_input, lengths, seq_axis=1)

        cell = _make_rnn_cell(rnn_spec)
        outputs, final_state = tf.nn.dynamic_rnn(
            cell, rnn_input, sequence_length=lengths, dtype=tf.float32)

        if rnn_spec.direction == "backward":
            outputs = tf.reverse_sequence(outputs, lengths, seq_axis=1)

        if rnn_spec.cell_type == "LSTM":
            final_state = final_state.h

    return outputs, final_state


class RecurrentEncoder(ModelPart, TemporalStatefulWithOutput):

    # pylint: disable=too-many-arguments
    def __init__(self,
                 name: str,
                 input_sequence: TemporalStateful,
                 rnn_layers: List[RNNSpecTuple],
                 add_residual: bool = False,
                 add_layer_norm: bool = False,
                 dropout_keep_prob: float = 1.0,
                 reuse: ModelPart = None,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None,
                 initializers: InitializerSpecs = None) -> None:
        """Create a new instance of a recurrent encoder.

        Arguments:
            name: ModelPart name.
            input_seqeunce: The input sequence for the encoder.
            rnn_size: The dimension of the RNN hidden state vector.
            rnn_cell: One of "GRU", "NematusGRU", "LSTM". Which kind of memory
                cell to use.
            rnn_direction: One of "forward", "backward", "bidirectional". In
                what order to process the input sequence. Note that choosing
                "bidirectional" will double the resulting vector dimension as
                well as the number of encoder parameters.
            add_residual: Add residual connections to the RNN layer output.
            add_layer_norm: Add layer normalization after each RNN layer.
            dropout_keep_prob: 1 - dropout probability.
            save_checkpoint: ModelPart save checkpoint file.
            load_checkpoint: ModelPart load checkpoint file.
        """
        check_argument_types()
        ModelPart.__init__(self, name, reuse, save_checkpoint, load_checkpoint,
                           initializers)
        TemporalStatefulWithOutput.__init__(self)

        self.input_sequence = input_sequence
        self.dropout_keep_prob = dropout_keep_prob
        self.rnn_specs = [_make_rnn_spec(*r) for r in rnn_layers]
        self.add_residual = add_residual
        self.add_layer_norm = add_layer_norm

        if self.dropout_keep_prob <= 0.0 or self.dropout_keep_prob > 1.0:
            raise ValueError("Dropout keep prob must be inside (0,1].")

        self._variable_scope.set_initializer(
            tf.random_normal_initializer(stddev=0.001))
    # pylint: enable=too-many-arguments

    @tensor
    def rnn_input(self) -> tf.Tensor:
        return dropout(self.input_sequence.temporal_states,
                       self.dropout_keep_prob, self.train_mode)

    @tensor
    def rnn(self) -> Tuple[tf.Tensor, tf.Tensor]:
        layer_input = self.rnn_input  # type: tf.Tensor
        # pylint: disable=unsubscriptable-object
        layer_final = self.rnn_input[:, -1]
        # pylint: enable=unsubscriptable-object

        for i, rnn_spec in enumerate(self.rnn_specs):
            with tf.variable_scope("rnn_{}_{}".format(i, rnn_spec.direction),
                                   reuse=tf.AUTO_REUSE):

                if self.add_layer_norm:
                    layer_input = layer_norm(layer_input)

                layer_output, layer_final_output = rnn_layer(
                    layer_input, self.input_sequence.lengths, rnn_spec)

                layer_output = dropout(
                    layer_output, self.dropout_keep_prob, self.train_mode)
                layer_final_output = dropout(
                    layer_final_output, self.dropout_keep_prob,
                    self.train_mode)

                in_dim = layer_input.get_shape()[-1]
{'fraction_non_alphanumeric': 0.08569482288828338, 'fraction_numerical': 0.0032697547683923707, 'mean_word_length': 2.4381733021077285, 'pattern_counts': {'":': 16, '<': 7, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5089134', 'n_tokens_mistral': 4544, 'n_tokens_neox': 4350, 'n_words': 1074}