<commit_before>import tensorflow as tf
import numpy as np

from variational import VariationalEncoder, VariationalDecoder
from data_reader import CHILD_EDGE, SIBLING_EDGE

class Model():
    def __init__(self, args, infer=False):
        self.args = args
        if infer:
            args.batch_size = 1
            args.max_ast_depth = 1

        if args.cell == 'lstm':
            args.cell_fn = tf.nn.rnn_cell.BasicLSTMCell(args.rnn_size, state_is_tuple=False)
        else:
            args.cell_fn = tf.nn.rnn_cell.BasicRNNCell(args.rnn_size)

        # setup the encoder
        self.encoder = VariationalEncoder(args)

        # sample from Normal(0,1) and reparameterize
        samples = tf.random_normal([args.batch_size, args.latent_size], 0., 1., dtype=tf.float32)
        self.psi = self.encoder.psi_mean + (self.encoder.psi_stdv * samples)

        # setup the decoder with psi as the initial state
        expansion_w = tf.get_variable('expansion_w', [args.latent_size, args.cell_fn.state_size])
        expansion_b = tf.get_variable('expansion_b', [args.cell_fn.state_size])
        self.initial_state = tf.matmul(self.psi, expansion_w) + expansion_b
        self.decoder = VariationalDecoder(args, initial_state=self.initial_state, infer=infer)

        # get the decoder outputs
        output = tf.reshape(tf.concat(1, self.decoder.outputs), [-1, args.cell_fn.output_size])
        logits = tf.matmul(output, self.decoder.projection_w) + self.decoder.projection_b
        self.probs = tf.nn.softmax(logits)

        # define losses
        self.targets = tf.placeholder(tf.int32, [args.batch_size, args.max_ast_depth])
        self.latent_loss = 0.5 * tf.reduce_sum(tf.square(self.encoder.psi_mean)
                                    + tf.square(self.encoder.psi_stdv)
                                    - tf.log(tf.square(self.encoder.psi_stdv)) - 1, 1)
        self.generation_loss = tf.nn.seq2seq.sequence_loss([logits],
                                    [tf.reshape(self.targets, [-1])],
                                    [tf.ones([args.batch_size * args.max_ast_depth])])
        self.cost = tf.reduce_mean(args.weight_loss * self.generation_loss + self.latent_loss)
        self.train_op = tf.train.AdamOptimizer(args.learning_rate).minimize(self.cost)

        var_params = [np.prod([dim.value for dim in var.get_shape()])
                            for var in tf.trainable_variables()]
        if not infer:
            print('Model parameters: {}'.format(np.sum(var_params)))

    def infer(self, sess, seqs, nodes, edges, input_vocab, target_vocab, psi=None):<commit_msg>Change in variational objective (cost) function<commit_after>import tensorflow as tf
import numpy as np

from variational import VariationalEncoder, VariationalDecoder
from data_reader import CHILD_EDGE, SIBLING_EDGE

class Model():
    def __init__(self, args, infer=False):
        self.args = args
        if infer:
            args.batch_size = 1
            args.max_ast_depth = 1

        if args.cell == 'lstm':
            args.cell_fn = tf.nn.rnn_cell.BasicLSTMCell(args.rnn_size, state_is_tuple=False)
        else:
            args.cell_fn = tf.nn.rnn_cell.BasicRNNCell(args.rnn_size)

        # setup the encoder
        self.encoder = VariationalEncoder(args)

        # sample from Normal(0,1) and reparameterize
        samples = tf.random_normal([args.batch_size, args.latent_size], 0., 1., dtype=tf.float32)
        self.psi = self.encoder.psi_mean + (self.encoder.psi_stdv * samples)

        # setup the decoder with psi as the initial state
        expansion_w = tf.get_variable('expansion_w', [args.latent_size, args.cell_fn.state_size])
        expansion_b = tf.get_variable('expansion_b', [args.cell_fn.state_size])
        self.initial_state = tf.matmul(self.psi, expansion_w) + expansion_b
        self.decoder = VariationalDecoder(args, initial_state=self.initial_state, infer=infer)

        # get the decoder outputs
        output = tf.reshape(tf.concat(1, self.decoder.outputs), [-1, args.cell_fn.output_size])
        logits = tf.matmul(output, self.decoder.projection_w) + self.decoder.projection_b
        self.probs = tf.nn.softmax(logits)

        # define losses
        self.targets = tf.placeholder(tf.int32, [args.batch_size, args.max_ast_depth])
        self.latent_loss = 0.5 * tf.reduce_sum(tf.square(self.encoder.psi_mean)
                                    + tf.square(self.encoder.psi_stdv)
                                    - tf.log(tf.square(self.encoder.psi_stdv)) - 1, 1)
        self.generation_loss = tf.nn.seq2seq.sequence_loss([logits],
                                    [tf.reshape(self.targets, [-1])],
                                    [tf.ones([args.batch_size * args.max_ast_depth])])
        self.cost = tf.reduce_mean(self.generation_loss + self.latent_loss/args.weight_loss)
        self.train_op = tf.train.AdamOptimizer(args.learning_rate).minimize(self.cost)

        var_params = [np.prod([dim.value for dim in var.get_shape()])
                            for var in tf.trainable_variables()]
        if not infer:
            print('Model parameters: {}'.format(np.sum(var_params)))

    def infer(self, sess, seqs, nodes, edges, input_vocab, target_vocab, psi=None):
{'fraction_non_alphanumeric': 0.11123809523809523, 'fraction_numerical': 0.006857142857142857, 'mean_word_length': 2.5123745819397993, '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': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6471484', 'n_tokens_mistral': 1717, 'n_tokens_neox': 1649, 'n_words': 303}