<commit_before>from typing import List, Optional
from typeguard import check_argument_types

import tensorflow as tf

from neuralmonkey.dataset import Dataset
from neuralmonkey.decorators import tensor
from neuralmonkey.model.model_part import ModelPart, FeedDict, InitializerSpecs
from neuralmonkey.model.stateful import Stateful, SpatialStatefulWithOutput


# pylint: disable=too-few-public-methods
class StatefulFiller(ModelPart, Stateful):
    """Placeholder class for stateful input.

    This model part is used to feed 1D tensors to the model. Optionally, it
    projects the states to given dimension.
    """

    def __init__(self,
                 name: str,
                 dimension: int,
                 data_id: str,
                 output_shape: int = None,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None,
                 initializers: InitializerSpecs = None) -> None:
        """Instantiate StatefulFiller.

        Arguments:
            name: Name of the model part.
            dimension: Dimensionality of the input.
            data_id: Series containing the numpy objects.
            output_shape: Dimension of optional state projection.
        """
        check_argument_types()
        ModelPart.__init__(
            self, name, save_checkpoint, load_checkpoint, initializers)

        self.data_id = data_id
        self.dimension = dimension
        self.output_shape = output_shape

        if self.dimension <= 0:
            raise ValueError("Input vector dimension must be positive.")
        if self.output_shape is not None and self.output_shape <= 0:
            raise ValueError("Output vector dimension must be positive.")

        with self.use_scope():
            self.vector = tf.placeholder(
                tf.float32, [None, self.dimension], "input_vector")

    @tensor
    def output(self) -> tf.Tensor:
        if self.output_shape is None or self.dimension == self.output_shape:
            return self.vector

        return tf.layers.dense(self.vector, self.output_shape)

    def feed_dict(self, dataset: Dataset, train: bool = False) -> FeedDict:
        fd = ModelPart.feed_dict(self, dataset, train)
        fd[self.vector] = dataset.get_series(self.data_id)
        return fd


class SpatialFiller(ModelPart, SpatialStatefulWithOutput):
    """Placeholder class for 3D numerical input.

    This model part is used to feed 3D tensors (e.g., pre-trained convolutional
    maps image captioning). Optionally, the states are projected to given size.
    """

    def __init__(self,
                 name: str,
                 input_shape: List[int],
                 data_id: str,
                 projection_dim: int = None,
                 save_checkpoint: Optional[str] = None,
                 load_checkpoint: Optional[str] = None,
                 initializers: InitializerSpecs = None) -> None:
        """Instantiate SpatialFiller.

        Args:
            name: Name of the model part.
            input_shape: Dimensionality of the input.
            data_id: Name of the data series with numpy objects.
            projection_dim: Optional, dimension of the states projection.
        """
        check_argument_types()
        ModelPart.__init__(
            self, name, save_checkpoint, load_checkpoint, initializers)

        self.data_id = data_id
        self.input_shape = input_shape
        self.projection_dim = projection_dim

        if len(self.input_shape) != 3:
            raise ValueError("The input shape should have 3 dimensions.")

        features_shape = [None] + self.input_shape  # type: ignore
        with self.use_scope():
            self.spatial_input = tf.placeholder(
                tf.float32, shape=features_shape, name="spatial_states")

    @tensor
    def output(self) -> tf.Tensor:
        return tf.reduce_mean(
            self.spatial_states, axis=[1, 2], name="average_image")

    @tensor
    def spatial_states(self) -> tf.Tensor:
        if self.projection_dim:
            return tf.layers.conv2d(
                self.spatial_input, filters=self.projection_dim,
                kernel_size=1, activation=None)
        return self.spatial_input

    @tensor
    def spatial_mask(self) -> tf.Tensor:
        return tf.ones(tf.shape(self.spatial_states)[:3])

    def feed_dict(self, dataset: Dataset, train: bool = False) -> FeedDict:
        fd = ModelPart.feed_dict(self, dataset, train)
        fd[self.spatial_input] = dataset.get_series(self.data_id)
        return fd<commit_msg>allow stateful filler to optionally do projection<commit_after>from typing import List
from typeguard import check_argument_types

import tensorflow as tf

from neuralmonkey.dataset import Dataset
from neuralmonkey.decorators import tensor
from neuralmonkey.model.model_part import ModelPart, FeedDict, InitializerSpecs
from neuralmonkey.model.stateful import Stateful, SpatialStatefulWithOutput


# pylint: disable=too-few-public-methods
class StatefulFiller(ModelPart, Stateful):
    """Placeholder class for stateful input.

    This model part is used to feed 1D tensors to the model. Optionally, it
    projects the states to given dimension.
    """

    def __init__(self,
                 name: str,
                 dimension: int,
                 data_id: str,
                 output_shape: int = None,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None,
                 initializers: InitializerSpecs = None) -> None:
        """Instantiate StatefulFiller.

        Arguments:
            name: Name of the model part.
            dimension: Dimensionality of the input.
            data_id: Series containing the numpy objects.
            output_shape: Dimension of optional state projection.
        """
        check_argument_types()
        ModelPart.__init__(
            self, name, save_checkpoint, load_checkpoint, initializers)

        self.data_id = data_id
        self.dimension = dimension
        self.output_shape = output_shape

        if self.dimension <= 0:
            raise ValueError("Input vector dimension must be positive.")
        if self.output_shape is not None and self.output_shape <= 0:
            raise ValueError("Output vector dimension must be positive.")

        with self.use_scope():
            self.vector = tf.placeholder(
                tf.float32, [None, self.dimension], "input_vector")

    @tensor
    def output(self) -> tf.Tensor:
        if self.output_shape is None or self.dimension == self.output_shape:
            return self.vector

        return tf.layers.dense(self.vector, self.output_shape)

    def feed_dict(self, dataset: Dataset, train: bool = False) -> FeedDict:
        fd = ModelPart.feed_dict(self, dataset, train)
        fd[self.vector] = dataset.get_series(self.data_id)
        return fd


class SpatialFiller(ModelPart, SpatialStatefulWithOutput):
    """Placeholder class for 3D numerical input.

    This model part is used to feed 3D tensors (e.g., pre-trained convolutional
    maps image captioning). Optionally, the states are projected to given size.
    """

    # pylint: disable=too-many-arguments
    def __init__(self,
                 name: str,
                 input_shape: List[int],
                 data_id: str,
                 projection_dim: int = None,
                 ff_hidden_dim: int = None,
                 save_checkpoint: str = None,
                 load_checkpoint: str = None,
                 initializers: InitializerSpecs = None) -> None:
        """Instantiate SpatialFiller.

        Args:
            name: Name of the model part.
            input_shape: Dimensionality of the input.
            data_id: Name of the data series with numpy objects.
            projection_dim: Optional, dimension of the states projection.
        """
        check_argument_types()
        ModelPart.__init__(
            self, name, save_checkpoint, load_checkpoint, initializers)

        self.data_id = data_id
        self.input_shape = input_shape
        self.projection_dim = projection_dim
        self.ff_hidden_dim = ff_hidden_dim

        if self.ff_hidden_dim is not None and self.projection_dim is None:
            raise ValueError(
                "projection_dim must be provided when using ff_hidden_dim")

        if len(self.input_shape) != 3:
            raise ValueError("The input shape should have 3 dimensions.")

        features_shape = [None] + self.input_shape  # type: ignore
        with self.use_scope():
            self.spatial_input = tf.placeholder(
                tf.float32, shape=features_shape, name="spatial_states")
    # pylint: enable=too-many-arguments

    @tensor
    def output(self) -> tf.Tensor:
        return tf.reduce_mean(
            self.spatial_states, axis=[1, 2], name="average_image")

    @tensor
    def spatial_states(self) -> tf.Tensor:
        if self.ff_hidden_dim:
            projected = tf.layers.conv2d(
                self.spatial_input, filters=self.ff_hidden_dim,
                kernel_size=1, activation=tf.nn.relu)
        else:
            projected = self.spatial_input

        if self.projection_dim:
            return tf.layers.conv2d(
                projected, filters=self.projection_dim,
                kernel_size=1, activation=None)

        assert projected == self.spatial_input
        return self.spatial_input

    @tensor
    def spatial_mask(self) -> tf.Tensor:
        return tf.ones(tf.shape(self.spatial_states)[:3])

    def feed_dict(self, dataset: Dataset, train: bool = False) -> FeedDict:
        fd = ModelPart.feed_dict(self, dataset, train)
        fd[self.spatial_input] = dataset.get_series(self.data_id)
        return fd
{'fraction_non_alphanumeric': 0.0793797494352023, 'fraction_numerical': 0.0034914766892585746, 'mean_word_length': 2.2648340596714718, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2359129', 'n_tokens_mistral': 2723, 'n_tokens_neox': 2597, 'n_words': 749}