<commit_before>from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import absolute_import
from builtins import *  # NOQA
from future import standard_library
standard_library.install_aliases()

import chainer
from chainer import functions as F
from chainer import links as L


class LinearBN(chainer.Chain):
    """Linear layer with BatchNormalization."""

    def __init__(self, in_size, out_size):
        super().__init__()
        with self.init_scope():
            self.linear = L.Linear(in_size, out_size)
            bn = L.BatchNormalization(out_size)
            bn.avg_var[:] = 1
            self.bn = bn

    def __call__(self, x):
        return self.bn(self.linear(x))


class MLPBN(chainer.Chain):
    """Multi-Layer Perceptron with BatchNormalization."""

    def __init__(self, in_size, out_size, hidden_sizes, normalize_input=True,
                 normalize_output=False):
        self.in_size = in_size
        self.out_size = out_size
        self.hidden_sizes = hidden_sizes
        self.normalize_input = normalize_input
        self.normalize_output = normalize_output

        super().__init__()
        with self.init_scope():
            if normalize_input:
                self.input_bn = L.BatchNormalization(in_size)
                self.input_bn.avg_var[:] = 1

            if hidden_sizes:
                hidden_layers = []
                hidden_layers.append(LinearBN(in_size, hidden_sizes[0]))
                for hin, hout in zip(hidden_sizes, hidden_sizes[1:]):
                    hidden_layers.append(LinearBN(hin, hout))
                self.hidden_layers = chainer.ChainList(*hidden_layers)
                self.output = L.Linear(hidden_sizes[-1], out_size)
            else:
                self.output = L.Linear(in_size, out_size)

            if normalize_output:
                self.output_bn = L.BatchNormalization(out_size)
                self.output_bn.avg_var[:] = 1

    def __call__(self, x):
        h = x
        assert (not chainer.config.train) or x.shape[0] > 1
        if self.normalize_input:
            h = self.input_bn(h)
        if self.hidden_sizes:
            for l in self.hidden_layers:
                h = F.relu(l(h))
        h = self.output(h)
        if self.normalize_output:
            h = self.output_bn(h)
        return h<commit_msg>Add arguments to MLPBN<commit_after>from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import absolute_import
from builtins import *  # NOQA
from future import standard_library
standard_library.install_aliases()

import chainer
from chainer import functions as F
from chainer import links as L

from chainerrl.initializers import LeCunNormal


class LinearBN(chainer.Chain):
    """Linear layer with BatchNormalization."""

    def __init__(self, in_size, out_size):
        super().__init__()
        with self.init_scope():
            self.linear = L.Linear(in_size, out_size)
            bn = L.BatchNormalization(out_size)
            bn.avg_var[:] = 1
            self.bn = bn

    def __call__(self, x):
        return self.bn(self.linear(x))


class MLPBN(chainer.Chain):
    """Multi-Layer Perceptron with BatchNormalization."""

    def __init__(self, in_size, out_size, hidden_sizes, normalize_input=True,
                 normalize_output=False, nonlinearity=F.relu, last_wscale=1):
        self.in_size = in_size
        self.out_size = out_size
        self.hidden_sizes = hidden_sizes
        self.normalize_input = normalize_input
        self.normalize_output = normalize_output
        self.nonlinearity = nonlinearity

        super().__init__()
        with self.init_scope():
            if normalize_input:
                self.input_bn = L.BatchNormalization(in_size)
                self.input_bn.avg_var[:] = 1

            if hidden_sizes:
                hidden_layers = []
                hidden_layers.append(LinearBN(in_size, hidden_sizes[0]))
                for hin, hout in zip(hidden_sizes, hidden_sizes[1:]):
                    hidden_layers.append(LinearBN(hin, hout))
                self.hidden_layers = chainer.ChainList(*hidden_layers)
                self.output = L.Linear(hidden_sizes[-1], out_size,
                                       initialW=LeCunNormal(last_wscale))
            else:
                self.output = L.Linear(in_size, out_size,
                                       initialW=LeCunNormal(last_wscale))

            if normalize_output:
                self.output_bn = L.BatchNormalization(out_size)
                self.output_bn.avg_var[:] = 1

    def __call__(self, x):
        h = x
        assert (not chainer.config.train) or x.shape[0] > 1
        if self.normalize_input:
            h = self.input_bn(h)
        if self.hidden_sizes:
            for l in self.hidden_layers:
                h = self.nonlinearity(l(h))
        h = self.output(h)
        if self.normalize_output:
            h = self.output_bn(h)
        return h
{'fraction_non_alphanumeric': 0.08680213059775103, 'fraction_numerical': 0.0033537186821858355, 'mean_word_length': 2.1588785046728973, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '3187360', 'n_tokens_mistral': 1564, 'n_tokens_neox': 1485, 'n_words': 326}