<filename>nlp/encode_blocks.py<gh_stars>10-100
import tensorflow as tf

from nlp.nn import get_var, layer_norm


def CNN_encode(seqs, filter_size=3, num_filters=None, scope=None,
               reuse=None, **kwargs):
    """
    Basic CNN encodes on the length dimension
    """
    seq_dim = seqs.get_shape().as_list()[-1]
    num_filters = num_filters if num_filters else seq_dim

    with tf.variable_scope(scope or 'conv_pool_%s' % filter_size, reuse=reuse):
        W = get_var('conv_W', [filter_size, seq_dim, num_filters])
        b = get_var('conv_bias', [num_filters])

        conv = tf.nn.conv1d(
            seqs,
            W,
            stride=1,
            padding='SAME',
            name='conv')

        # apply nonlinearity
        return tf.nn.relu(tf.nn.bias_add(conv, b))


def MH_Att_encode(queries,
                  keys,
                  num_units=None,
                  num_heads=8,
                  dropout_keep_rate=1.0,
                  causality=False,
                  scope='MultiHead_Attention_Block',
                  reuse=None,
                  residual=False,
                  normalize_output=True,
                  **kwargs):
    """Applies multihead attention.

    Args:
      queries: A 3d tensor with shape of [N, T_q, C_q].
      keys: A 3d tensor with shape of [N, T_k, C_k].
      num_units: A scalar. Attention size.
      dropout_rate: A floating point number.
      is_training: Boolean. Controller of mechanism for dropout.
      causality: Boolean. If true, units that reference the future are masked.
      num_heads: An int. Number of heads.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns
      A 3d tensor with shape of (N, T_q, C)
    """
    if num_units is None or residual:
        num_units = queries.get_shape().as_list()[-1]
    with tf.variable_scope(scope, reuse=reuse):
        # Set the fall back option for num_units

        # Linear projections
        Q = tf.layers.dense(queries, num_units, activation=tf.nn.relu)  # (N, T_q, C)
        K = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)
        V = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)

        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)  # (h*N, T_q, C/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)

        # Multiplication
        outputs = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1]))  # (h*N, T_q, T_k)

        # Scale
        outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)

        # Key Masking
        key_masks = tf.sign(tf.abs(tf.reduce_sum(keys, axis=-1)))  # (N, T_k)
        key_masks = tf.tile(key_masks, [num_heads, 1])  # (h*N, T_k)
        key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1])  # (h*N, T_q, T_k)

        paddings = tf.ones_like(outputs) * (-2 ** 32 + 1)
        outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Causality = Future blinding
        if causality:
            diag_vals = tf.ones_like(outputs[0, :, :])  # (T_q, T_k)
            tril = tf.contrib.linalg.LinearOperatorTriL(diag_vals).to_dense()  # (T_q, T_k)
            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1])  # (h*N, T_q, T_k)

            paddings = tf.ones_like(masks) * (-2 ** 32 + 1)
            outputs = tf.where(tf.equal(masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Activation
        outputs = tf.nn.softmax(outputs)  # (h*N, T_q, T_k)

        # Query Masking
        query_masks = tf.sign(tf.abs(tf.reduce_sum(queries, axis=-1)))  # (N, T_q)
        query_masks = tf.tile(query_masks, [num_heads, 1])  # (h*N, T_q)
        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]])  # (h*N, T_q, T_k)
        outputs *= query_masks  # broadcasting. (N, T_q, C)

        # Dropouts
        outputs = tf.nn.dropout(outputs, keep_prob=dropout_keep_rate)

        # Weighted sum
        outputs = tf.matmul(outputs, V_)  # ( h*N, T_q, C/h)

        # Restore shape
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)  # (N, T_q, C)

        if residual:
            # Residual connection
            outputs += queries

        if normalize_output:
            # Normalize
            outputs = layer_norm(outputs)  # (N, T_q, C)

    return outputs

<filename>nlp/encode_blocks.py<gh_stars>10-100
import tensorflow as tf

from nlp.nn import get_var, layer_norm


def CNN_encode(seqs, filter_size=3, num_filters=None, scope=None,
               reuse=None, **kwargs):
    """
    Basic CNN encodes on the length dimension
    """
    seq_dim = seqs.get_shape().as_list()[-1]
    num_filters = num_filters if num_filters else seq_dim

    with tf.variable_scope(scope or 'conv_pool_%s' % filter_size, reuse=reuse):
        W = get_var('conv_W', [filter_size, seq_dim, num_filters])
        b = get_var('conv_bias', [num_filters])

        conv = tf.nn.conv1d(
            seqs,
            W,
            stride=1,
            padding='SAME',
            name='conv')

        # apply nonlinearity
        return tf.nn.relu(tf.nn.bias_add(conv, b))


def MH_Att_encode(queries,
                  keys,
                  num_units=None,
                  num_heads=8,
                  dropout_keep_rate=1.0,
                  causality=False,
                  scope='MultiHead_Attention_Block',
                  reuse=None,
                  residual=False,
                  normalize_output=True,
                  **kwargs):
    """Applies multihead attention.

    Args:
      queries: A 3d tensor with shape of [N, T_q, C_q].
      keys: A 3d tensor with shape of [N, T_k, C_k].
      num_units: A scalar. Attention size.
      dropout_rate: A floating point number.
      is_training: Boolean. Controller of mechanism for dropout.
      causality: Boolean. If true, units that reference the future are masked.
      num_heads: An int. Number of heads.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns
      A 3d tensor with shape of (N, T_q, C)
    """
    if num_units is None or residual:
        num_units = queries.get_shape().as_list()[-1]
    with tf.variable_scope(scope, reuse=reuse):
        # Set the fall back option for num_units

        # Linear projections
        Q = tf.layers.dense(queries, num_units, activation=tf.nn.relu)  # (N, T_q, C)
        K = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)
        V = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)

        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)  # (h*N, T_q, C/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)

        # Multiplication
        outputs = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1]))  # (h*N, T_q, T_k)

        # Scale
        outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)

        # Key Masking
        key_masks = tf.sign(tf.abs(tf.reduce_sum(keys, axis=-1)))  # (N, T_k)
        key_masks = tf.tile(key_masks, [num_heads, 1])  # (h*N, T_k)
        key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1])  # (h*N, T_q, T_k)

        paddings = tf.ones_like(outputs) * (-2 ** 32 + 1)
        outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Causality = Future blinding
        if causality:
            diag_vals = tf.ones_like(outputs[0, :, :])  # (T_q, T_k)
            tril = tf.contrib.linalg.LinearOperatorTriL(diag_vals).to_dense()  # (T_q, T_k)
            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1])  # (h*N, T_q, T_k)

            paddings = tf.ones_like(masks) * (-2 ** 32 + 1)
            outputs = tf.where(tf.equal(masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Activation
        outputs = tf.nn.softmax(outputs)  # (h*N, T_q, T_k)

        # Query Masking
        query_masks = tf.sign(tf.abs(tf.reduce_sum(queries, axis=-1)))  # (N, T_q)
        query_masks = tf.tile(query_masks, [num_heads, 1])  # (h*N, T_q)
        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]])  # (h*N, T_q, T_k)
        outputs *= query_masks  # broadcasting. (N, T_q, C)

        # Dropouts
        outputs = tf.nn.dropout(outputs, keep_prob=dropout_keep_rate)

        # Weighted sum
        outputs = tf.matmul(outputs, V_)  # ( h*N, T_q, C/h)

        # Restore shape
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)  # (N, T_q, C)

        if residual:
            # Residual connection
            outputs += queries

        if normalize_output:
            # Normalize
            outputs = layer_norm(outputs)  # (N, T_q, C)

    return outputs

Basic CNN encodes on the length dimension # apply nonlinearity Applies multihead attention.      Args:       queries: A 3d tensor with shape of [N, T_q, C_q].       keys: A 3d tensor with shape of [N, T_k, C_k].       num_units: A scalar. Attention size.       dropout_rate: A floating point number.       is_training: Boolean. Controller of mechanism for dropout.       causality: Boolean. If true, units that reference the future are masked.       num_heads: An int. Number of heads.       scope: Optional scope for `variable_scope`.       reuse: Boolean, whether to reuse the weights of a previous layer         by the same name.      Returns       A 3d tensor with shape of (N, T_q, C) # Set the fall back option for num_units # Linear projections # (N, T_q, C) # (N, T_k, C) # (N, T_k, C) # Split and concat # (h*N, T_q, C/h) # (h*N, T_k, C/h) # (h*N, T_k, C/h) # Multiplication # (h*N, T_q, T_k) # Scale # Key Masking # (N, T_k) # (h*N, T_k) # (h*N, T_q, T_k) # (h*N, T_q, T_k) # Causality = Future blinding # (T_q, T_k) # (T_q, T_k) # (h*N, T_q, T_k) # (h*N, T_q, T_k) # Activation # (h*N, T_q, T_k) # Query Masking # (N, T_q) # (h*N, T_q) # (h*N, T_q, T_k) # broadcasting. (N, T_q, C) # Dropouts # Weighted sum # ( h*N, T_q, C/h) # Restore shape # (N, T_q, C) # Residual connection # Normalize # (N, T_q, C)