<commit_before>from typing import Set, cast

import tensorflow as tf
from typeguard import check_argument_types

from neuralmonkey.model.stateful import TemporalStatefulWithOutput
from neuralmonkey.model.model_part import ModelPart, FeedDict
from neuralmonkey.nn.utils import dropout
from neuralmonkey.dataset import Dataset
from neuralmonkey.decorators import tensor
from neuralmonkey.attention.base_attention import (
    get_attention_states, get_attention_mask, Attendable)


class AttentiveEncoder(ModelPart, TemporalStatefulWithOutput):
    """An encoder with attention over the input and a fixed-dimension output.

    Based on "A Structured Self-attentive Sentence Embedding",
    https://arxiv.org/abs/1703.03130.

    The encoder combines a sequence of vectors into a fixed-size matrix where
    each row of the matrix is computed using a different attention head. This
    matrix is exposed as the ``temporal_states`` property (the time dimension
    corresponds to the different attention heads). The ``output`` property
    provides a flattened or projected representation of this matrix.
    """

    # pylint: disable=too-many-arguments
    def __init__(self,
                 name: str,
                 input_sequence: Attendable,
                 hidden_size: int,
                 num_heads: int,
                 output_size: int = None,
                 state_proj_size: int = None,
                 dropout_keep_prob: float = 1.0,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None) -> None:
        """Initialize an instance of the encoder."""
        check_argument_types()
        ModelPart.__init__(self, name, save_checkpoint, load_checkpoint)
        TemporalStatefulWithOutput.__init__(self)

        self.input_sequence = input_sequence
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.output_size = output_size
        self.state_proj_size = state_proj_size
        self.dropout_keep_prob = dropout_keep_prob

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

        with self.use_scope():
            self.train_mode = tf.placeholder(tf.bool, [], "train_mode")
    # pylint: enable=too-many-arguments

    @tensor
    def temporal_states(self) -> tf.Tensor:
        states = dropout(get_attention_states(self.input_sequence),
                         self.dropout_keep_prob,
                         self.train_mode)
        mask = get_attention_mask(self.input_sequence)

        hidden = tf.layers.dense(states, units=self.hidden_size,
                                 activation=tf.tanh, use_bias=False,
                                 name="S1")
        energies = tf.layers.dense(hidden, units=self.num_heads,
                                   use_bias=False, name="S2")
        # shape: [batch_size, max_time, num_heads]
        weights = tf.nn.softmax(energies, dim=1)
        if mask is not None:
            weights *= tf.expand_dims(mask, -1)
            weights /= tf.reduce_sum(weights, axis=1, keep_dims=True) + 1e-8

        self._weights = weights

        if self.state_proj_size is not None:
            states = tf.layers.dense(states, units=self.state_proj_size,
                                     name="state_projection")

        return tf.matmul(a=weights, b=states, transpose_a=True)

    @tensor
    def temporal_mask(self) -> tf.Tensor:
        return tf.ones(tf.shape(self.temporal_states)[:2], tf.float32)

    @tensor
    def output(self) -> tf.Tensor:
        # pylint: disable=no-member
        state_size = self.temporal_states.get_shape()[2].value
        output = tf.reshape(self.temporal_states,
                            [-1, self.num_heads * state_size])
        if self.output_size is not None:
            output = tf.layers.dense(output, self.output_size,
                                     name="output_projection")

        return output

    @tensor
    def attention_weights(self) -> tf.Tensor:
        _ = self.temporal_states
        return self._weights

    def get_dependencies(self) -> Set[ModelPart]:
        deps = ModelPart.get_dependencies(self)

        # feed only if needed
        if isinstance(self.input_sequence, ModelPart):
            feedable = cast(ModelPart, self.input_sequence)
            deps |= feedable.get_dependencies()

        return deps

    def feed_dict(self, dataset: Dataset, train: bool = False) -> FeedDict:
        return {self.train_mode: train}<commit_msg>Add attention_states property to attentive encoder<commit_after>from typing import Set, cast

import tensorflow as tf
from typeguard import check_argument_types

from neuralmonkey.model.stateful import TemporalStatefulWithOutput
from neuralmonkey.model.model_part import ModelPart, FeedDict
from neuralmonkey.nn.utils import dropout
from neuralmonkey.dataset import Dataset
from neuralmonkey.decorators import tensor
from neuralmonkey.attention.base_attention import (
    get_attention_states, get_attention_mask, Attendable)


class AttentiveEncoder(ModelPart, TemporalStatefulWithOutput):
    """An encoder with attention over the input and a fixed-dimension output.

    Based on "A Structured Self-attentive Sentence Embedding",
    https://arxiv.org/abs/1703.03130.

    The encoder combines a sequence of vectors into a fixed-size matrix where
    each row of the matrix is computed using a different attention head. This
    matrix is exposed as the ``temporal_states`` property (the time dimension
    corresponds to the different attention heads). The ``output`` property
    provides a flattened and, optionally, projected representation of this
    matrix.
    """

    # pylint: disable=too-many-arguments
    def __init__(self,
                 name: str,
                 input_sequence: Attendable,
                 hidden_size: int,
                 num_heads: int,
                 output_size: int = None,
                 state_proj_size: int = None,
                 dropout_keep_prob: float = 1.0,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None) -> None:
        """Initialize an instance of the encoder."""
        check_argument_types()
        ModelPart.__init__(self, name, save_checkpoint, load_checkpoint)
        TemporalStatefulWithOutput.__init__(self)

        self.input_sequence = input_sequence
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.output_size = output_size
        self.state_proj_size = state_proj_size
        self.dropout_keep_prob = dropout_keep_prob

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

        with self.use_scope():
            self.train_mode = tf.placeholder(tf.bool, [], "train_mode")
    # pylint: enable=too-many-arguments

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

    @tensor
    def attention_weights(self) -> tf.Tensor:
        mask = get_attention_mask(self.input_sequence)
        hidden = tf.layers.dense(self.attention_states, units=self.hidden_size,
                                 activation=tf.tanh, use_bias=False,
                                 name="S1")
        energies = tf.layers.dense(hidden, units=self.num_heads,
                                   use_bias=False, name="S2")
        # shape: [batch_size, max_time, num_heads]
        weights = tf.nn.softmax(energies, dim=1)
        if mask is not None:
            weights *= tf.expand_dims(mask, -1)
            weights /= tf.reduce_sum(weights, axis=1, keep_dims=True) + 1e-8

        return weights

    @tensor
    def temporal_states(self) -> tf.Tensor:
        states = self.attention_states
        if self.state_proj_size is not None:
            states = tf.layers.dense(states, units=self.state_proj_size,
                                     name="state_projection")

        return tf.matmul(a=self.attention_weights, b=states, transpose_a=True)

    @tensor
    def temporal_mask(self) -> tf.Tensor:
        return tf.ones(tf.shape(self.temporal_states)[:2], tf.float32)

    @tensor
    def output(self) -> tf.Tensor:
        # pylint: disable=no-member
        state_size = self.temporal_states.get_shape()[2].value
        output = tf.reshape(self.temporal_states,
                            [-1, self.num_heads * state_size])
        if self.output_size is not None:
            output = tf.layers.dense(output, self.output_size,
                                     name="output_projection")

        return output

    def get_dependencies(self) -> Set[ModelPart]:
        deps = ModelPart.get_dependencies(self)

        # feed only if needed
        if isinstance(self.input_sequence, ModelPart):
            feedable = cast(ModelPart, self.input_sequence)
            deps |= feedable.get_dependencies()

        return deps

    def feed_dict(self, dataset: Dataset, train: bool = False) -> FeedDict:
        return {self.train_mode: train}
{'fraction_non_alphanumeric': 0.08319774378999892, 'fraction_numerical': 0.006291354810716998, 'mean_word_length': 2.3735821441639224, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 20, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5376353', 'n_tokens_mistral': 2677, 'n_tokens_neox': 2528, 'n_words': 632}