<commit_before>from keras.models import Model
from keras.layers import Input, Convolution2D, Activation, Flatten, Dense, Reshape, Lambda
from keras.layers.advanced_activations import PReLU
from keras.engine.topology import merge
from keras.optimizers import Nadam
import keras.backend.tensorflow_backend as K

from eva.layers.gated_cnn import GatedCNN, GatedCNNs
from eva.layers.masked_convolution2d import MaskedConvolution2D

def GatedPixelCNN(input_shape, filters, blocks, latent=0, build=True):
    width, height, channels = input_shape

    # TODO: Make it scalable to any amount of palettes.
    input_img = Input(shape=input_shape, name=str(channels)+'_channels_256_palette')

    latent_vector = None
    if latent != 0:
        latent_vector = Input(shape=(10,), name='latent_vector')

    model = GatedCNNs(filters, blocks, latent_vector)(*GatedCNN(filters, latent_vector)(input_img))

    for _ in range(2):
        model = Convolution2D(filters, 1, 1, border_mode='valid')(model)
        model = PReLU()(model)

    if channels == 1:
        outs = Convolution2D(1, 1, 1, activation='sigmoid', border_mode='valid')(model)
    else:
        out = Convolution2D(256*channels, 1, 1, border_mode='valid', name='channels_mult_palette')(model)
        out = Reshape((input_shape[0], input_shape[1], 256, input_shape[2]), name='palette_channels')(out)

        outs = [None] * channels
        for i in range(channels):
            outs[i] = Lambda(lambda x: x[:, :, :, :, 0], name='channel'+str(i)+'_extract')(out)
            outs[i] = Reshape((input_shape[0] * input_shape[1], 256), name='hw_palette'+str(i))(outs[i])
            outs[i] = Activation('softmax', name='channel'+str(i))(outs[i])

    if build:
        model = Model(input=[input_img, latent_vector] if latent != 0 else input_img, output=outs)
        model.compile(optimizer=Nadam(), loss='binary_crossentropy' if channels == 1 else 'sparse_categorical_crossentropy')

    return model
<commit_msg>Fix channels output; improve formatting<commit_after>from keras.models import Model
from keras.layers import Input, Convolution2D, Activation, Flatten, Dense, Reshape, Lambda
from keras.layers.advanced_activations import PReLU
from keras.engine.topology import merge
from keras.optimizers import Nadam
import keras.backend.tensorflow_backend as K

from eva.layers.gated_cnn import GatedCNN, GatedCNNs
from eva.layers.masked_convolution2d import MaskedConvolution2D

def GatedPixelCNN(input_shape, filters, blocks, latent=None, build=True):
    width, height, channels = input_shape
    palette = 256

    # TODO: Make it scalable to any amount of palettes.
    input_img = Input(shape=input_shape, name=str(channels)+'_channels_256_palette')

    latent_vector = None
    if latent is not None:
        latent_vector = Input(shape=(latent,), name='latent_vector')

    model = GatedCNNs(filters, blocks, latent_vector)(*GatedCNN(filters, latent_vector)(input_img))

    for _ in range(2):
        model = Convolution2D(filters, 1, 1, border_mode='valid')(model)
        model = PReLU()(model)

    if channels == 1:
        outs = Convolution2D(1, 1, 1, activation='sigmoid', border_mode='valid')(model)
    else:
        out = Convolution2D(palette*channels, 1, 1, border_mode='valid', name='channels_mult_palette')(model)
        out = Reshape((height, width, palette, channels), name='palette_channels')(out)

        outs = [None] * channels
        for i in range(channels):
            outs[i] = Lambda(lambda x: x[:, :, :, :, i], name='channel'+str(i)+'_extract')(out)
            outs[i] = Reshape((height * width, palette), name='hw_palette'+str(i))(outs[i])
            outs[i] = Activation('softmax', name='channel'+str(i))(outs[i])

    if build:
        model = Model(input=[input_img, latent_vector] if latent is not None else input_img, output=outs)
        model.compile(optimizer=Nadam(), loss='binary_crossentropy' if channels == 1 else 'sparse_categorical_crossentropy')

    return model

{'fraction_non_alphanumeric': 0.1246544357878864, 'fraction_numerical': 0.015330485046494095, 'mean_word_length': 4.4520547945205475, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '4484335', 'n_tokens_mistral': 1359, 'n_tokens_neox': 1268, 'n_words': 304}