Upload modules.py

modules.py

@@ -0,0 +1,233 @@
+import tensorflow as tf
+
+def embedding_lookup(lookup_table, x):
+    return tf.compat.v1.nn.embedding_lookup(lookup_table, x)
+
+
+def normal_embedding_lookup(x, n_token, d_embed, d_proj, initializer,
+                            proj_initializer, scope='normal_embed', **kwargs):
+    emb_scale = d_proj ** 0.5
+    with tf.compat.v1.variable_scope(scope):
+        lookup_table = tf.compat.v1.get_variable('lookup_table', [n_token, d_embed], initializer=initializer)
+        y = embedding_lookup(lookup_table, x)
+        if d_proj != d_embed:
+            proj_W = tf.compat.v1.get_variable('proj_W', [d_embed, d_proj], initializer=proj_initializer)
+            y = tf.einsum('ibe,ed->ibd', y, proj_W)
+        else:
+            proj_W = None
+        ret_params = [lookup_table, proj_W]
+    y *= emb_scale
+    return y, ret_params
+
+
+def normal_softmax(hidden, target, n_token, params, scope='normal_softmax', **kwargs):
+    def _logit(x, W, b, proj):
+        y = x
+        if proj is not None:
+            y = tf.einsum('ibd,ed->ibe', y, proj)
+        return tf.einsum('ibd,nd->ibn', y, W) + b
+
+    params_W, params_projs = params[0], params[1]
+
+    with tf.compat.v1.variable_scope(scope):
+        softmax_b = tf.compat.v1.get_variable('bias', [n_token], initializer=tf.zeros_initializer())
+        output = _logit(hidden, params_W, softmax_b, params_projs)
+        nll = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target, logits=output)
+    return nll, output
+
+
+def positional_embedding(pos_seq, inv_freq, bsz=None):
+    sinusoid_inp = tf.einsum('i,j->ij', pos_seq, inv_freq)
+    pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)
+    if bsz is not None:
+        return tf.tile(pos_emb[:, None, :], [1, bsz, 1])
+    else:
+        return pos_emb[:, None, :]
+
+
+def positionwise_FF(inp, d_model, d_inner, dropout, kernel_initializer,
+                    scope='ff', is_training=True):
+    output = inp
+    with tf.compat.v1.variable_scope(scope):
+        output = tf.keras.layers.Dense(d_inner, activation=tf.nn.relu, 
+                                       kernel_initializer=kernel_initializer, name='layer_1')(inp)
+        output = tf.keras.layers.Dropout(dropout, name='drop_1')(output, training=is_training)
+        output = tf.keras.layers.Dense(d_model, activation=tf.nn.relu, 
+                                       kernel_initializer=kernel_initializer, name='layer_2')(output)
+        output = tf.keras.layers.Dropout(dropout, name='drop_2')(output, training=is_training)
+        output = tf.keras.layers.LayerNormalization(axis=-1)(output + inp)
+    return output
+
+
+def _create_mask(qlen, mlen, same_length=False):
+    attn_mask = tf.ones([qlen, qlen])
+    mask_u = tf.linalg.band_part(attn_mask, 0, -1)
+    mask_dia = tf.linalg.band_part(attn_mask, 0, 0)
+    attn_mask_pad = tf.zeros([qlen, mlen])
+    ret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)
+    if same_length:
+        mask_l = tf.matrix_band_part(attn_mask, -1, 0)
+        ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)
+    return ret
+
+
+def _cache_mem(curr_out, prev_mem, mem_len=None):
+    if mem_len is None or prev_mem is None:
+        new_mem = curr_out
+    elif mem_len == 0:
+        return prev_mem
+    else:
+        new_mem = tf.concat([prev_mem, curr_out], 0)[-mem_len:]
+    return tf.stop_gradient(new_mem)
+
+
+def rel_shift(x):
+    x_size = tf.shape(x)
+    x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]])
+    x = tf.reshape(x, [x_size[1] + 1, x_size[0], x_size[2], x_size[3]])
+    x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
+    x = tf.reshape(x, x_size)
+    return x
+
+
+def rel_multihead_attn(w, r, r_w_bias, r_r_bias, attn_mask, mems, d_model,
+                       n_head, d_head, dropout, dropatt, is_training,
+                       kernel_initializer, scope='rel_attn'):
+    scale = 1 / (d_head ** 0.5)
+    with tf.compat.v1.variable_scope(scope):
+        qlen = tf.shape(w)[0]
+        rlen = tf.shape(r)[0]
+        bsz = tf.shape(w)[1]
+
+        cat = tf.concat([mems, w], 0) if mems is not None and mems.shape.ndims > 1 else w
+
+        w_heads = tf.keras.layers.Dense(3 * n_head * d_head, use_bias=False, 
+                                        kernel_initializer=kernel_initializer, name='qkv')(cat)
+        r_head_k = tf.keras.layers.Dense(n_head * d_head, use_bias=False,
+                                         kernel_initializer=kernel_initializer, name='r')(r)
+        
+        w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, -1)
+        w_head_q = w_head_q[-qlen:]
+
+        klen = tf.shape(w_head_k)[0]
+
+        w_head_q = tf.reshape(w_head_q, [qlen, bsz, n_head, d_head])
+        w_head_k = tf.reshape(w_head_k, [klen, bsz, n_head, d_head])
+        w_head_v = tf.reshape(w_head_v, [klen, bsz, n_head, d_head])
+
+        r_head_k = tf.reshape(r_head_k, [rlen, n_head, d_head])
+
+        rw_head_q = w_head_q + r_w_bias
+        rr_head_q = w_head_q + r_r_bias
+
+        AC = tf.einsum('ibnd,jbnd->ijbn', rw_head_q, w_head_k)
+        BD = tf.einsum('ibnd,jnd->ijbn', rr_head_q, r_head_k)
+        BD = rel_shift(BD)
+
+        attn_score = (AC + BD) * scale
+        attn_mask_t = attn_mask[:, :, None, None]
+        attn_score = attn_score * (1 - attn_mask_t) - 1e30 * attn_mask_t
+
+        attn_prob = tf.nn.softmax(attn_score, 1)
+        attn_prob = tf.keras.layers.Dropout(dropatt)(attn_prob, training=is_training)
+
+        attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, w_head_v)
+        size_t = tf.shape(attn_vec)
+        attn_vec = tf.reshape(attn_vec, [size_t[0], size_t[1], n_head * d_head])
+
+        attn_out = tf.keras.layers.Dense(d_model, use_bias=False, 
+                                         kernel_initializer=kernel_initializer, name='o')(attn_vec)
+        attn_out = tf.keras.layers.Dropout(dropout)(attn_out, training=is_training)
+        output = tf.keras.layers.LayerNormalization(axis=-1)(attn_out + w)
+        return output
+
+
+def transformer(dec_inp, target, mems, n_token, n_layer, d_model, d_embed,
+                n_head, d_head, d_inner, dropout, dropatt,
+                initializer, is_training, proj_initializer=None,
+                mem_len=None, cutoffs=[], div_val=1, tie_projs=[],
+                same_length=False, clamp_len=-1,
+                input_perms=None, target_perms=None, head_target=None,
+                untie_r=False, proj_same_dim=True,
+                scope='transformer'):
+    """
+    cutoffs: a list of python int. Cutoffs for adaptive softmax.
+    tie_projs: a list of python bools. Whether to tie the projections.
+    perms: a list of tensors. Each tensor should of size [len, bsz, bin_size].
+        Only used in the adaptive setting.
+    """
+    new_mems = []
+    with tf.compat.v1.variable_scope(scope):
+        if untie_r:
+            r_w_bias = tf.compat.v1.get_variable('r_w_bias', [n_layer, n_head, d_head], initializer=initializer)
+            r_r_bias = tf.compat.v1.get_variable('r_r_bias', [n_layer, n_head, d_head], initializer=initializer)
+        else:
+            r_w_bias = tf.compat.v1.get_variable('r_w_bias', [n_head, d_head], initializer=initializer)
+            r_r_bias = tf.compat.v1.get_variable('r_r_bias', [n_head, d_head], initializer=initializer)
+
+        qlen = tf.shape(dec_inp)[0]
+        mlen = tf.shape(mems[0])[0] if mems is not None else 0
+        klen = qlen + mlen
+
+        if proj_initializer is None:
+            proj_initializer = initializer
+
+        embeddings, shared_params = normal_embedding_lookup(
+            x=dec_inp,
+            n_token=n_token,
+            d_embed=d_embed,
+            d_proj=d_model,
+            initializer=initializer,
+            proj_initializer=proj_initializer)
+        
+        attn_mask = _create_mask(qlen, mlen, same_length)
+        
+        pos_seq = tf.range(klen - 1, -1, -1.0)
+        if clamp_len > 0:
+            pos_seq = tf.minimum(pos_seq, clamp_len)
+        inv_freq = 1 / (10000 ** (tf.range(0, d_model, 2.0) / d_model))
+        pos_emb = positional_embedding(pos_seq, inv_freq)
+
+        output = tf.keras.layers.Dropout(rate=dropout)(embeddings, training=is_training)
+        pos_emb = tf.keras.layers.Dropout(rate=dropout)(pos_emb, training=is_training)
+
+        if mems is None:
+            mems = [None] * n_layer
+
+        for i in range(n_layer):
+            # cache new mems
+            new_mems.append(_cache_mem(output, mems[i], mem_len))
+
+            with tf.compat.v1.variable_scope('layer_{}'.format(i)):
+                output = rel_multihead_attn(
+                    w=output,
+                    r=pos_emb,
+                    r_w_bias=r_w_bias if not untie_r else r_w_bias[i],
+                    r_r_bias=r_r_bias if not untie_r else r_r_bias[i],
+                    attn_mask=attn_mask,
+                    mems=mems[i],
+                    d_model=d_model,
+                    n_head=n_head,
+                    d_head=d_head,
+                    dropout=dropout,
+                    dropatt=dropatt,
+                    is_training=is_training,
+                    kernel_initializer=initializer)
+
+                output = positionwise_FF(
+                    inp=output,
+                    d_model=d_model,
+                    d_inner=d_inner,
+                    dropout=dropout,
+                    kernel_initializer=initializer,
+                    is_training=is_training)
+
+        output = tf.keras.layers.Dropout(dropout)(output, training=is_training)
+
+        loss, logits = normal_softmax(
+            hidden=output,
+            target=target,
+            n_token=n_token,
+            params=shared_params)
+
+        return loss, logits, new_mems