<commit_before>"""

Variable notations:
s: synaptic connection weight
c: synaptic connection probability
b: leak bias

"""
import os
import shutil
import numpy as np
import tensorflow as tf
from itertools import permutations

from .truenorth import TrueNorthChip, NEURON_DEST_OUT
from .util import maybe_mkdir

# Default types for tensors
TF_DTYPE = tf.float32
NP_DTYPE = np.float32
NP_DTYPE_INT = np.int32

TF_MODEL_NAME = 'model.ckpt'


def normal_ccdf(x, mu, sigma2):
    """Normal CCDF"""
    # Check for degenerate distributions when sigma2 == 0
    # if x >= mu, n = 0
    # if x < mu, n = 1
    # sigma2_le_0 = tf.less_equal(sigma2, 0.)
    # x_gte_mu = tf.greater_equal(x, mu)
    # x_lt_mu = tf.less(x, mu)

    # Never divide by zero, instead the logic below handles degenerate distribution cases
    # sigma2 = tf.cond(sigma2_le_0, lambda: tf.ones_like(sigma2), lambda: sigma2)

    p = (1. - 0.5 * (1. + tf.erf((x - mu) / tf.sqrt(2. * sigma2))))
    # p = tf.cond(tf.logical_and(sigma2_le_0, x_gte_mu), lambda: tf.zeros_like(p), lambda: p)
    # p = tf.cond(tf.logical_and(sigma2_le_0, x_lt_mu), lambda: tf.ones_like(p), lambda: p)
    return p


def leak_bias(shape, core_name):
    """Create a bias variable, which is implemented as a leak"""
    return tf.Variable(tf.zeros(shape, dtype=TF_DTYPE), name='%s_leak_bias' % core_name)


def synapse_connection(shape, core_name):
    """Create a variable for the synaptic connection probability"""
    return tf.Variable(tf.random_uniform(shape, 0., 1., dtype=TF_DTYPE,
                                         seed=np.random.randint(0, np.iinfo(np.int32).max)),
                       name='%s_synapse_connection' % core_name)


def synapse_weight(shape, core_name, weights=np.array([-2, -1, 1, 2])):
    """Create a constant for the synaptic connection weights, which remain fixed during training"""
    # Evenly distribute the 4 axon types by first choosing the axon types
    axon_types = np.repeat(np.arange(4, dtype=NP_DTYPE_INT), shape[0] / 4)[:shape[0]]
    axon_types = np.random.permutation(axon_types)

    # Then compute all possible permutations from the weight choices
    axon_weight_perms = np.array(list(permutations(weights)), dtype=NP_DTYPE_INT)

    # And index into the array of permutations
    axon_weight_idx = np.tile(np.arange(len(axon_weight_perms)), np.ceil(shape[1] / len(axon_weight_perms)))[:shape[1]]
    axon_weight_idx = np.random.permutation(axon_weight_idx)

    axon_weights = axon_weight_perms[axon_weight_idx]

    # Create the full synaptic weight matrix from the types and weights
    s = axon_weights[:, axon_types].T.astype(NP_DTYPE)

    return tf.Variable(s, name='%s_synapse_weight' % core_name, trainable=False), \
           tf.Variable(axon_types, name='%s_axon_type' % core_name, trainable=False), \
           tf.Variable(axon_weights, name='%s_axon_weight' % core_name, trainable=False)


def connect_cores(input, output_dim, name):
    """Connect two cores given the inputs, synaptic weights, and output dimension.
    Inputs can be output from a previous core or spike inputs"""
    input_dim = int(input.get_shape()[1])

    s, axon_types, axon_weights = synapse_weight((input_dim, output_dim), name)
    b = leak_bias([output_dim], name)
    c = synapse_connection([input_dim, output_dim], name)

    xc = tf.reshape(input, (-1, input_dim, 1)) * c
    mu = b + tf.reduce_sum(xc * s, 1)
    sigma2 = tf.reduce_sum(xc * (1. - xc) * tf.pow(s, 2), 1)

    # Output is proba that each neuron fires
    x0 = tf.zeros_like(mu)
    output = normal_ccdf(x0, mu, sigma2)

    return output, b, c, axon_types, axon_weights, s


def configure_core(core, b, c, axon_types, axon_weights, dest_core, dest_axon):
    """Configure a TrueNorth core, given the corresponding parameters of the shadow network"""
    num_axons, num_neurons = c.shape

    # Set the axon types and weights, only up to the neurons for this core
    core.sigma[:num_neurons, :] = np.sign(axon_weights)
    core.s[:num_neurons, :] = np.abs(axon_weights)

    # Use the bias as a leak
    core.sigma_lambda[:num_neurons] = np.sign(b)
    core.lambda_[:num_neurons] = np.abs(b)
    core.sigma_lambda[core.sigma_lambda == 0] = 1

    # Most of these correspond to default values and shouldn't change
    core.b[:] = False
    core.c_lambda[:] = False
    core.epsilon[:] = False
    core.alpha[:] = 0
    core.beta[:] = 0
    core.TM[:] = 0
    core.gamma[:] = 0
    core.kappa[:] = True
    core.sigma_VR[:] = 1
    core.VR[:] = 0
    core.V[:] = 0

    # Set the axon types and weights
    core.axon_type[:num_axons] = axon_types

    # Set the synaptic connections
    core.w[:num_axons, :num_neurons] = c > np.random.uniform(0, 1, size=c.shape)

    core.dest_core[:num_neurons] = dest_core
    core.dest_axon[:num_neurons] = dest_axon

    return core


class TrueShadow(object):
    """The probabilistic interpretation of a spiking network"""

    def __init__(self):
        """The model must either be loaded or trained to be tested or deployed"""
        self.sess = None

    def load(self, dirpath, model_name=TF_MODEL_NAME):
        """Load the model from a checkpoint dir"""
        model_fname = os.path.join(dirpath, model_name)
        saver = tf.train.Saver()
        sess = tf.Session()
        ckpt = tf.train.get_checkpoint_state(dirpath)

        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            self.sess = sess
        else:
            print('Unable to load from checkpoint:', model_fname)

        return

    def train(self, train_data, validation_data=None,
              batch_size=100, num_epochs=10, sample_weights=None,
              checkpoint_steps=1, report_steps=1, deploy_steps=0,
              dirpath=None, model_name=TF_MODEL_NAME, dropout=0.5):
        """Train the network and optionally save checkpoints"""

        if dirpath:
            ckpt_dir = maybe_mkdir(dirpath)
            model_fname = os.path.join(ckpt_dir, model_name)

        # Op to initialize all variables which should have been created in a subclass initialization
        init_op = tf.global_variables_initializer()

        # Create a saver
        saver = tf.train.Saver()

        # Create the session and save it for future use. Only other place sessions are created is in .load()
        sess = tf.Session()

        sess.run(init_op)
        self.sess = sess

        train_size = len(train_data)
        tf.train.start_queue_runners(sess=sess)

        if sample_weights is None:
            sample_weights = np.ones(len(train_data))

        results = []
        for epoch_i in range(1, num_epochs + 1):
            for batch_i in range(train_size // batch_size):
                batch_xs, batch_ys = train_data.next_batch(batch_size)

                sess.run(self.optimizer,
                         feed_dict={
                             self.x: batch_xs,
                             self.y: batch_ys,
                             self.sample_weights: sample_weights,
                             self.keep_prob: dropout,
                         })

                # Clip the bias variables to [-255, 255] and connection probas to [0, 1]
                for (core_b, core_c, _, _, _) in self.core_params.values():
                    sess.run(tf.assign(core_b, tf.clip_by_value(core_b, -255., 255.)))
                    sess.run(tf.assign(core_c, tf.clip_by_value(core_c, 0., 1.)))

            report = ''
            report_values = [epoch_i]

            if report_steps > 0 and (epoch_i % report_steps == 0 or epoch_i == num_epochs):
                train_acc = self.test(train_data, batch_size=batch_size)
                test_acc = self.test(validation_data, batch_size=batch_size)
                report += '[ Train  ACC: %.4f ][ Test  ACC: %.4f ]'
                report_values.append(train_acc)
                report_values.append(test_acc)

            if deploy_steps > 0 and (epoch_i % deploy_steps == 0 or epoch_i == num_epochs):<commit_msg>Fixed param for sample weights<commit_after>"""

Variable notations:
s: synaptic connection weight
c: synaptic connection probability
b: leak bias

"""
import os
import shutil
import numpy as np
import tensorflow as tf
from itertools import permutations

from .truenorth import TrueNorthChip, NEURON_DEST_OUT
from .util import maybe_mkdir

# Default types for tensors
TF_DTYPE = tf.float32
NP_DTYPE = np.float32
NP_DTYPE_INT = np.int32

TF_MODEL_NAME = 'model.ckpt'


def normal_ccdf(x, mu, sigma2):
    """Normal CCDF"""
    # Check for degenerate distributions when sigma2 == 0
    # if x >= mu, n = 0
    # if x < mu, n = 1
    # sigma2_le_0 = tf.less_equal(sigma2, 0.)
    # x_gte_mu = tf.greater_equal(x, mu)
    # x_lt_mu = tf.less(x, mu)

    # Never divide by zero, instead the logic below handles degenerate distribution cases
    # sigma2 = tf.cond(sigma2_le_0, lambda: tf.ones_like(sigma2), lambda: sigma2)

    p = (1. - 0.5 * (1. + tf.erf((x - mu) / tf.sqrt(2. * sigma2))))
    # p = tf.cond(tf.logical_and(sigma2_le_0, x_gte_mu), lambda: tf.zeros_like(p), lambda: p)
    # p = tf.cond(tf.logical_and(sigma2_le_0, x_lt_mu), lambda: tf.ones_like(p), lambda: p)
    return p


def leak_bias(shape, core_name):
    """Create a bias variable, which is implemented as a leak"""
    return tf.Variable(tf.zeros(shape, dtype=TF_DTYPE), name='%s_leak_bias' % core_name)


def synapse_connection(shape, core_name):
    """Create a variable for the synaptic connection probability"""
    return tf.Variable(tf.random_uniform(shape, 0., 1., dtype=TF_DTYPE,
                                         seed=np.random.randint(0, np.iinfo(np.int32).max)),
                       name='%s_synapse_connection' % core_name)


def synapse_weight(shape, core_name, weights=np.array([-2, -1, 1, 2])):
    """Create a constant for the synaptic connection weights, which remain fixed during training"""
    # Evenly distribute the 4 axon types by first choosing the axon types
    axon_types = np.repeat(np.arange(4, dtype=NP_DTYPE_INT), shape[0] / 4)[:shape[0]]
    axon_types = np.random.permutation(axon_types)

    # Then compute all possible permutations from the weight choices
    axon_weight_perms = np.array(list(permutations(weights)), dtype=NP_DTYPE_INT)

    # And index into the array of permutations
    axon_weight_idx = np.tile(np.arange(len(axon_weight_perms)), np.ceil(shape[1] / len(axon_weight_perms)))[:shape[1]]
    axon_weight_idx = np.random.permutation(axon_weight_idx)

    axon_weights = axon_weight_perms[axon_weight_idx]

    # Create the full synaptic weight matrix from the types and weights
    s = axon_weights[:, axon_types].T.astype(NP_DTYPE)

    return tf.Variable(s, name='%s_synapse_weight' % core_name, trainable=False), \
           tf.Variable(axon_types, name='%s_axon_type' % core_name, trainable=False), \
           tf.Variable(axon_weights, name='%s_axon_weight' % core_name, trainable=False)


def connect_cores(input, output_dim, name):
    """Connect two cores given the inputs, synaptic weights, and output dimension.
    Inputs can be output from a previous core or spike inputs"""
    input_dim = int(input.get_shape()[1])

    s, axon_types, axon_weights = synapse_weight((input_dim, output_dim), name)
    b = leak_bias([output_dim], name)
    c = synapse_connection([input_dim, output_dim], name)

    xc = tf.reshape(input, (-1, input_dim, 1)) * c
    mu = b + tf.reduce_sum(xc * s, 1)
    sigma2 = tf.reduce_sum(xc * (1. - xc) * tf.pow(s, 2), 1)

    # Output is proba that each neuron fires
    x0 = tf.zeros_like(mu)
    output = normal_ccdf(x0, mu, sigma2)

    return output, b, c, axon_types, axon_weights, s


def configure_core(core, b, c, axon_types, axon_weights, dest_core, dest_axon):
    """Configure a TrueNorth core, given the corresponding parameters of the shadow network"""
    num_axons, num_neurons = c.shape

    # Set the axon types and weights, only up to the neurons for this core
    core.sigma[:num_neurons, :] = np.sign(axon_weights)
    core.s[:num_neurons, :] = np.abs(axon_weights)

    # Use the bias as a leak
    core.sigma_lambda[:num_neurons] = np.sign(b)
    core.lambda_[:num_neurons] = np.abs(b)
    core.sigma_lambda[core.sigma_lambda == 0] = 1

    # Most of these correspond to default values and shouldn't change
    core.b[:] = False
    core.c_lambda[:] = False
    core.epsilon[:] = False
    core.alpha[:] = 0
    core.beta[:] = 0
    core.TM[:] = 0
    core.gamma[:] = 0
    core.kappa[:] = True
    core.sigma_VR[:] = 1
    core.VR[:] = 0
    core.V[:] = 0

    # Set the axon types and weights
    core.axon_type[:num_axons] = axon_types

    # Set the synaptic connections
    core.w[:num_axons, :num_neurons] = c > np.random.uniform(0, 1, size=c.shape)

    core.dest_core[:num_neurons] = dest_core
    core.dest_axon[:num_neurons] = dest_axon

    return core


class TrueShadow(object):
    """The probabilistic interpretation of a spiking network"""

    def __init__(self):
        """The model must either be loaded or trained to be tested or deployed"""
        self.sess = None

    def load(self, dirpath, model_name=TF_MODEL_NAME):
        """Load the model from a checkpoint dir"""
        model_fname = os.path.join(dirpath, model_name)
        saver = tf.train.Saver()
        sess = tf.Session()
        ckpt = tf.train.get_checkpoint_state(dirpath)

        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            self.sess = sess
        else:
            print('Unable to load from checkpoint:', model_fname)

        return

    def train(self, train_data, validation_data=None,
              batch_size=100, num_epochs=10, class_weight=None,
              checkpoint_steps=1, report_steps=1, deploy_steps=0,
              dirpath=None, model_name=TF_MODEL_NAME, dropout=0.5):
        """Train the network and optionally save checkpoints"""

        if dirpath:
            ckpt_dir = maybe_mkdir(dirpath)
            model_fname = os.path.join(ckpt_dir, model_name)

        # Op to initialize all variables which should have been created in a subclass initialization
        init_op = tf.global_variables_initializer()

        # Create a saver
        saver = tf.train.Saver()

        # Create the session and save it for future use. Only other place sessions are created is in .load()
        sess = tf.Session()

        sess.run(init_op)
        self.sess = sess

        train_size = len(train_data)
        tf.train.start_queue_runners(sess=sess)

        results = []
        for epoch_i in range(1, num_epochs + 1):
            for batch_i in range(train_size // batch_size):
                batch_xs, batch_ys = train_data.next_batch(batch_size)

                if sample_weights is None:
                    sample_weights = np.ones(len(batch_ys))
                else:
                    sample_weights = (batch_ys * class_weight).sum(axis=1)

                sess.run(self.optimizer,
                         feed_dict={
                             self.x: batch_xs,
                             self.y: batch_ys,
                             self.sample_weights: sample_weights,
                             self.keep_prob: dropout,
                         })

                # Clip the bias variables to [-255, 255] and connection probas to [0, 1]
                for (core_b, core_c, _, _, _) in self.core_params.values():
                    sess.run(tf.assign(core_b, tf.clip_by_value(core_b, -255., 255.)))
                    sess.run(tf.assign(core_c, tf.clip_by_value(core_c, 0., 1.)))

            report = ''
            report_values = [epoch_i]

            if report_steps > 0 and (epoch_i % report_steps == 0 or epoch_i == num_epochs):
                train_acc = self.test(train_data, batch_size=batch_size)
                test_acc = self.test(validation_data, batch_size=batch_size)
                report += '[ Train  ACC: %.4f ][ Test  ACC: %.4f ]'
                report_values.append(train_acc)
                report_values.append(test_acc)

            if deploy_steps > 0 and (epoch_i % deploy_steps == 0 or epoch_i == num_epochs):
{'fraction_non_alphanumeric': 0.10919326899282356, 'fraction_numerical': 0.012682504330611235, 'mean_word_length': 2.7212246777163904, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2612051', 'n_tokens_mistral': 5568, 'n_tokens_neox': 5183, 'n_words': 1359}