def model()

in src/model.py [0:0]

• 56 lines of code
• 8 McCabe index (conditional complexity)

def model(hparams, X, Y=None, past=None, scope='model', reuse=False):
    with tf.variable_scope(scope, reuse=reuse):
        results = {}
        batch, sequence = shape_list(X)

        if hparams.bert:
            M = tf.greater(tf.random.uniform([batch, sequence]), hparams.bert_mask_prob)
            M = tf.cast(M, tf.float32)

        wpe = tf.get_variable('wpe', [hparams.n_ctx, hparams.n_embd],
                             initializer=tf.random_normal_initializer(stddev=0.01))
        wte = tf.get_variable('wte', [hparams.n_vocab, hparams.n_embd],
                             initializer=tf.random_normal_initializer(stddev=0.02))
        wtet = tf.get_variable('wtet', [hparams.n_vocab, hparams.n_embd],
                             initializer=tf.random_normal_initializer(stddev=0.0))
        past_length = 0 if past is None else tf.shape(past)[-2]

        h = tf.gather(wte, X)

        if hparams.bert:
            h = h * tf.expand_dims(M, 2)
        else:
            sos = tf.get_variable('sos', [hparams.n_embd],
                                 initializer=tf.random_normal_initializer(stddev=0.02))
            sos_tok = tf.ones([batch, 1, hparams.n_embd], dtype=tf.float32) * sos
            h = tf.concat([sos_tok, h[:,:-1,:]], axis=1)

        h += tf.gather(wpe, positions_for(X, past_length))

        # Transformer
        presents = []
        pasts = tf.unstack(past, axis=1) if past is not None else [None] * hparams.n_layer
        assert len(pasts) == hparams.n_layer
        for layer, past in enumerate(pasts):
            h, present = block(h, 'h%d' % layer, past=past, hparams=hparams)
            presents.append(present)
        results['present'] = tf.stack(presents, axis=1)
        h = norm(h, 'ln_f')

        # Generative loss.  Do tokens <n predict token n?
        h_flat = tf.reshape(h, [batch*sequence, hparams.n_embd])
        gen_logits = tf.matmul(h_flat, wtet, transpose_b=True)
        gen_logits = tf.reshape(gen_logits, [batch, sequence, hparams.n_vocab])
        results['gen_logits'] = gen_logits

        gen_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=gen_logits, labels=X)
        if hparams.bert:
            IM = 1.0 - M
            gen_losses = tf.reduce_sum(gen_losses * IM, axis=1) / tf.reduce_sum(IM, axis=1)
            results['gen_loss'] = tf.reduce_mean(gen_losses)
        else:
            results['gen_loss'] = tf.reduce_mean(gen_losses)

        # Classification loss.
        with tf.variable_scope('clf', reuse=reuse):
            classes = shape_list(Y)[1]
            if hparams.clf:
                wclf = tf.get_variable('wclf', [classes, hparams.n_embd],
                                      initializer=tf.random_normal_initializer(stddev=0.0))
            else:
                wclf = tf.zeros([classes, hparams.n_embd], dtype=tf.float32)

        h = tf.reduce_mean(h, axis=1)  # average pool over sequence
        clf_logits = tf.matmul(h, wclf, transpose_b=True)
        clf_losses = tf.nn.softmax_cross_entropy_with_logits_v2(logits=clf_logits, labels=Y)
        results['clf_loss'] = tf.reduce_mean(clf_losses)

        correct = tf.equal(tf.argmax(clf_logits, -1), tf.argmax(Y, -1))
        results['accuracy'] = tf.reduce_mean(tf.cast(correct, tf.float32)) * 100.0

        return results