def lm_loss(hidden, target, n_token, d_model, initializer, lookup_table=None,
tie_weight=False, bi_data=True, use_tpu=False):
'''doc.'''
with tf.variable_scope('lm_loss'):
if tie_weight:
assert lookup_table is not None, \
'lookup_table cannot be None for tie_weight'
softmax_w = lookup_table
else:
softmax_w = tf.get_variable(
'weight', [n_token, d_model],
dtype=hidden.dtype, initializer=initializer)
softmax_b = tf.get_variable(
'bias', [n_token], dtype=hidden.dtype,
initializer=tf.zeros_initializer())
logits = tf.einsum('ibd,nd->ibn', hidden, softmax_w) + softmax_b
preds = tf.argmax(logits, axis=-1)
if use_tpu:
one_hot_target = tf.one_hot(target, n_token, dtype=logits.dtype)
loss = -tf.reduce_sum(
tf.nn.log_softmax(logits) * one_hot_target, -1)
else:
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=target, logits=logits)
return loss, preds
def __init__(self,
is_training,
input_tensor,
label_ids,
sample_weight=None,
scope='mrc',
name='',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
seq_length = input_tensor.shape.as_list()[-2]
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[2, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[2],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
output_layer = tf.reshape(output_layer, [-1, hidden_size])
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits, [-1, seq_length, 2])
logits = tf.transpose(logits, [0, 2, 1])
probs = tf.nn.softmax(logits, axis=-1, name='probs')
self.probs[name] = probs
start_one_hot_labels = tf.one_hot(label_ids[:, 0],
depth=seq_length,
dtype=tf.float32)
end_one_hot_labels = tf.one_hot(label_ids[:, 1],
depth=seq_length,
dtype=tf.float32)
start_log_probs = tf.nn.log_softmax(logits[:, 0, :], axis=-1)
end_log_probs = tf.nn.log_softmax(logits[:, 1, :], axis=-1)
per_example_loss = (
-0.5 * tf.reduce_sum(start_one_hot_labels * start_log_probs,
axis=-1) - 0.5 *
tf.reduce_sum(end_one_hot_labels * end_log_probs, axis=-1))
if sample_weight is not None:
per_example_loss *= sample_weight
self.total_loss = tf.reduce_mean(per_example_loss)
self.losses[name] = per_example_loss
start_preds = tf.expand_dims(tf.argmax(logits[:, 0, :], axis=-1),
axis=-1)
end_preds = tf.expand_dims(tf.argmax(logits[:, 1, :], axis=-1),
axis=-1)
self.preds[name] = tf.concat([start_preds, end_preds], axis=-1)
def __init__(self,
is_training,
input_tensor,
input_mask,
label_ids,
label_size=2,
sample_weight=None,
scope='cls/sequence',
name='',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
batch_size = tf.shape(input_tensor)[0]
seq_length = input_tensor.shape.as_list()[-2]
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
output_layer = tf.reshape(output_layer, [-1, hidden_size])
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits, [-1, seq_length, label_size])
self.preds[name] = tf.argmax(logits, axis=-1)
self.probs[name] = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(label_ids,
depth=label_size,
dtype=tf.float32)
per_token_losses = -tf.reduce_mean(one_hot_labels * log_probs,
axis=-1)
input_mask = tf.concat([
tf.zeros((batch_size, 1), dtype=tf.float32),
tf.cast(input_mask[:, 2:], dtype=tf.float32),
tf.zeros((batch_size, 1), dtype=tf.float32)
],
axis=-1)
per_token_losses *= input_mask
per_example_loss = tf.reduce_mean(per_token_losses, axis=-1)
if sample_weight is not None:
per_example_loss *= tf.cast(sample_weight, dtype=tf.float32)
self.losses[name] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
def __init__(self,
is_training,
input_tensor,
label_ids,
label_size=2,
sample_weight=None,
scope='cls/seq_relationship',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
self.preds['preds'] = tf.argmax(logits, axis=-1)
self.probs['probs'] = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(
label_ids, depth=label_size, dtype=tf.float32)
per_example_loss = - tf.reduce_sum(
one_hot_labels * log_probs, axis=-1)
if sample_weight is not None:
per_example_loss = tf.cast(
sample_weight, dtype=tf.float32) * per_example_loss
thresh = kwargs.get('tsa_thresh')
if thresh is not None:
assert isinstance(thresh, float), (
'`tsa_thresh` must be a float between 0 and 1.')
uncertainty = tf.reduce_sum(self.probs['probs'] * tf.log(
self.probs['probs']), axis=-1)
uncertainty /= tf.log(1 / label_size)
per_example_loss = tf.cast(
tf.greater(uncertainty, thresh), dtype=tf.float32) * \
per_example_loss
self.losses['losses'] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
def _forward(target_ids, target_mask, target_max_seq_length):
with tf.variable_scope('decoder'):
# shared embedding
dec = tf.nn.embedding_lookup(embedding_table, target_ids)
dec *= hidden_size ** 0.5 # scale
dec += positional_encoding(dec, target_max_seq_length)
dec = util.dropout(dec, dropout_rate)
# blocks
for i in range(num_blocks):
with tf.variable_scope('block_%s' % i):
# masked self-attention
dec = multihead_attention(
queries=dec,
keys=dec,
values=dec,
key_masks=target_mask,
num_heads=num_attention_heads,
dropout_rate=dropout_rate,
training=is_training,
causality=True,
scope='masked_self_attention')
# vanilla attention
dec = multihead_attention(
queries=dec,
keys=memory,
values=memory,
key_masks=source_mask,
num_heads=num_attention_heads,
dropout_rate=dropout_rate,
training=is_training,
causality=False,
scope='vanilla_attention')
# feed forward
dec = ff(
dec, num_units=[4 * hidden_size, hidden_size])
# final linear projection (embedding weights are shared)
with tf.variable_scope('cls'):
output_bias = tf.get_variable(
'output_bias', shape=[vocab_size],
initializer=tf.zeros_initializer())
dec = tf.reshape(dec, [-1, hidden_size])
logits = tf.matmul(dec, embedding_table, transpose_b=True)
logits = tf.reshape(
logits, [-1, target_max_seq_length, vocab_size])
logits = tf.nn.bias_add(logits, output_bias)
return logits
def __init__(self,
is_training,
input_tensor,
input_mask,
label_ids,
label_size=5,
sample_weight=None,
scope='cls/sequence',
name='',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
seq_length = input_tensor.shape.as_list()[-2]
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
output_layer = tf.reshape(output_layer, [-1, hidden_size])
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits, [-1, seq_length, label_size])
with tf.variable_scope('crf'):
input_length = tf.reduce_sum(input_mask, axis=-1)
per_example_loss, transition_matrix = \
contrib.crf.crf_log_likelihood(
inputs=logits,
tag_indices=label_ids,
sequence_lengths=input_length)
per_example_loss = -per_example_loss
if sample_weight is not None:
per_example_loss *= tf.cast(sample_weight,
dtype=tf.float32)
self.total_loss = tf.reduce_mean(per_example_loss)
self.losses[name] = per_example_loss
self.preds[name] = tf.argmax(logits, axis=-1)
self.probs['logits'] = logits
self.probs['transition_matrix'] = transition_matrix
def __init__(self,
is_training,
input_tensor,
label_ids,
label_size=2,
sample_weight=None,
label_weight=None,
scope='cls/seq_relationship',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
probs = tf.nn.sigmoid(logits, name='probs')
self.probs['probs'] = probs
self.preds['preds'] = tf.greater(probs, 0.5)
per_example_loss = tf.nn.sigmoid_cross_entropy_with_logits(
logits=logits,
labels=tf.cast(label_ids, dtype=tf.float32))
if label_weight is not None:
label_weight = tf.constant(label_weight, dtype=tf.float32)
label_weight = tf.reshape(label_weight, [1, label_size])
per_example_loss *= label_weight
per_example_loss = tf.reduce_mean(per_example_loss, axis=-1)
if sample_weight is not None:
per_example_loss *= sample_weight
self.losses['losses'] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
def __init__(self,
is_training,
input_tensor,
label_ids,
label_size=2,
sample_weight=None,
scope='cls/seq_relationship',
name='',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
self.preds[name] = tf.argmax(logits, axis=-1)
self.probs[name] = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(label_ids,
depth=label_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs,
axis=-1)
if sample_weight is not None:
per_example_loss = tf.cast(sample_weight,
dtype=tf.float32) * per_example_loss
self.losses[name] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
def _get_logits(pooled_output, hidden_size, scope, trainable):
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(
bert_config.initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
logits = tf.matmul(pooled_output,
output_weights,
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
return logits
def ln(inputs, epsilon = 1e-8, scope='ln'):
'''Applies layer normalization. See https://arxiv.org/abs/1607.06450.
inputs: A tensor with 2 or more dimensions, where the first dimension has `batch_size`.
epsilon: A floating number. A very small number for preventing ZeroDivision Error.
scope: Optional scope for `variable_scope`.
Returns:
A tensor with the same shape and data dtype as `inputs`.
'''
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
inputs_shape = inputs.get_shape()
params_shape = inputs_shape[-1:]
mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)
beta= tf.get_variable('beta', params_shape, initializer=tf.zeros_initializer())
gamma = tf.get_variable('gamma', params_shape, initializer=tf.ones_initializer())
normalized = (inputs - mean) / ( (variance + epsilon) ** (.5) )
outputs = gamma * normalized + beta
return outputs
def _cls_fcn(self,
prev_output,
label_size,
hidden_size=768,
initializer_range=0.02,
dtype=tf.float32,
trainable=True):
with tf.variable_scope('output'):
cls_output_weights = tf.get_variable(
'output_weights', [hidden_size, label_size],
initializer=tf.truncated_normal_initializer(
stddev=initializer_range),
dtype=dtype,
trainable=trainable)
cls_output_bias = tf.get_variable(
'output_bias', [label_size],
initializer=tf.zeros_initializer(),
dtype=dtype,
trainable=trainable)
cls_logits = tf.matmul(prev_output[:, 0, :], cls_output_weights)
cls_output = tf.nn.bias_add(cls_logits, cls_output_bias)
return cls_output
def _get_generator_output(self, inputs, sample_weight, generator):
'''Masked language modeling softmax layer.'''
def gather_indexes(sequence_tensor, positions):
sequence_shape = util.get_shape_list(sequence_tensor, 3)
batch_size = sequence_shape[0]
seq_length = sequence_shape[1]
width = sequence_shape[2]
flat_offsets = tf.reshape(
tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
flat_positions = tf.reshape(positions + flat_offsets, [-1])
flat_sequence_tensor = tf.reshape(sequence_tensor,
[batch_size * seq_length, width])
output_tensor = tf.gather(flat_sequence_tensor, flat_positions)
return output_tensor
input_tensor = gather_indexes(generator.get_sequence_output(),
inputs.masked_lm_positions)
with tf.variable_scope('generator_predictions'):
input_tensor = tf.layers.dense(
input_tensor,
units=self.config.embedding_size,
activation=util.get_activation(self.bert_config.hidden_act),
kernel_initializer=util.create_initializer(
self.bert_config.initializer_range))
input_tensor = util.layer_norm(input_tensor)
output_bias = tf.get_variable('output_bias',
shape=[self.bert_config.vocab_size],
initializer=tf.zeros_initializer())
logits = tf.matmul(input_tensor,
generator.get_embedding_table(),
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
probs = tf.nn.softmax(logits, axis=-1, name='MLM_probs')
preds = tf.argmax(logits, axis=-1)
log_probs = tf.nn.log_softmax(logits, axis=-1)
label_ids = tf.reshape(inputs.masked_lm_ids, [-1])
masked_lm_weights = inputs.masked_lm_weights
if sample_weight is not None:
sample_weight = tf.expand_dims(tf.cast(sample_weight,
dtype=tf.float32),
axis=-1)
masked_lm_weights *= sample_weight
label_weights = tf.reshape(masked_lm_weights, [-1])
one_hot_labels = tf.one_hot(label_ids,
depth=self.bert_config.vocab_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(log_probs * one_hot_labels,
axis=[-1])
per_example_loss = label_weights * per_example_loss
numerator = tf.reduce_sum(per_example_loss)
denominator = tf.reduce_sum(label_weights) + 1e-6
loss = numerator / denominator
MLMOutput = collections.namedtuple(
'MLMOutput',
['logits', 'probs', 'loss', 'per_example_loss', 'preds'])
return MLMOutput(logits=logits,
probs=probs,
per_example_loss=per_example_loss,
loss=loss,
preds=preds)
def __init__(self,
bert_config,
is_training,
input_tensor,
sa_mask,
label_ids,
sample_weight=None,
scope='sanet',
alpha=0.5,
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
shape = util.get_shape_list(input_tensor)
batch_size = shape[0]
seq_length = shape[1]
hidden_size = shape[2]
sa_mask = tf.reshape(sa_mask, [batch_size, seq_length, seq_length])
with tf.variable_scope(scope):
with tf.variable_scope('sentence_attention'):
(sa_output, _) = self.attention_layer(
from_tensor=input_tensor,
to_tensor=input_tensor,
attention_mask=sa_mask,
num_attention_heads=bert_config.num_attention_heads,
size_per_head=\
hidden_size // bert_config.num_attention_heads,
attention_probs_dropout_prob=\
bert_config.hidden_dropout_prob,
initializer_range=bert_config.initializer_range,
do_return_2d_tensor=False,
batch_size=batch_size,
from_max_seq_length=seq_length,
to_max_seq_length=seq_length,
trainable=trainable)
with tf.variable_scope('cls/mrc'):
output_weights = tf.get_variable(
'output_weights',
shape=[2, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[2],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = alpha * sa_output + (1 - alpha) * input_tensor
output_layer = tf.reshape(output_layer, [-1, hidden_size])
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits, [-1, seq_length, 2])
logits = tf.transpose(logits, [0, 2, 1])
probs = tf.nn.softmax(logits, axis=-1, name='probs')
self.probs['probs'] = probs
self.preds['preds'] = tf.argmax(logits, axis=-1)
start_one_hot_labels = tf.one_hot(label_ids[:, 0],
depth=seq_length,
dtype=tf.float32)
end_one_hot_labels = tf.one_hot(label_ids[:, 1],
depth=seq_length,
dtype=tf.float32)
start_log_probs = tf.nn.log_softmax(logits[:, 0, :], axis=-1)
end_log_probs = tf.nn.log_softmax(logits[:, 1, :], axis=-1)
per_example_loss = (
-0.5 * tf.reduce_sum(start_one_hot_labels * start_log_probs,
axis=-1) - 0.5 *
tf.reduce_sum(end_one_hot_labels * end_log_probs, axis=-1))
if sample_weight is not None:
per_example_loss *= sample_weight
self.total_loss = tf.reduce_mean(per_example_loss)
self.losses['losses'] = per_example_loss
def __init__(self,
bert_config,
is_training,
encoder,
masked_lm_positions,
masked_lm_ids,
masked_lm_weights,
next_sentence_labels,
sample_weight=None,
scope_lm='cls/predictions',
scope_cls='cls/seq_relationship',
trainable=True,
use_nsp_loss=True,
**kwargs):
super(BERTDecoder, self).__init__(**kwargs)
def gather_indexes(sequence_tensor, positions):
sequence_shape = util.get_shape_list(sequence_tensor, 3)
batch_size = sequence_shape[0]
seq_length = sequence_shape[1]
width = sequence_shape[2]
flat_offsets = tf.reshape(
tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
flat_positions = tf.reshape(positions + flat_offsets, [-1])
flat_sequence_tensor = tf.reshape(sequence_tensor,
[batch_size * seq_length, width])
output_tensor = tf.gather(flat_sequence_tensor, flat_positions)
return output_tensor
scalar_losses = []
# masked language modeling
input_tensor = gather_indexes(encoder.get_sequence_output(),
masked_lm_positions)
with tf.variable_scope(scope_lm):
with tf.variable_scope('transform'):
input_tensor = tf.layers.dense(
input_tensor,
units=bert_config.hidden_size,
activation=util.get_activation(bert_config.hidden_act),
kernel_initializer=util.create_initializer(
bert_config.initializer_range))
input_tensor = util.layer_norm(input_tensor)
output_bias = tf.get_variable('output_bias',
shape=[bert_config.vocab_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
logits = tf.matmul(input_tensor,
encoder.get_embedding_table(),
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
probs = tf.nn.softmax(logits, axis=-1, name='MLM_probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
label_ids = tf.reshape(masked_lm_ids, [-1])
if sample_weight is not None:
sample_weight = tf.expand_dims(tf.cast(sample_weight,
dtype=tf.float32),
axis=-1)
masked_lm_weights *= sample_weight
label_weights = tf.reshape(masked_lm_weights, [-1])
one_hot_labels = tf.one_hot(label_ids,
depth=bert_config.vocab_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(log_probs * one_hot_labels,
axis=[-1])
per_example_loss = label_weights * per_example_loss
numerator = tf.reduce_sum(per_example_loss)
denominator = tf.reduce_sum(label_weights) + 1e-5
loss = numerator / denominator
scalar_losses.append(loss)
self.losses['MLM_losses'] = per_example_loss
self.preds['MLM_preds'] = tf.argmax(probs, axis=-1)
# next sentence prediction
with tf.variable_scope(scope_cls):
output_weights = tf.get_variable(
'output_weights',
shape=[2, bert_config.hidden_size],
initializer=util.create_initializer(
bert_config.initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[2],
initializer=tf.zeros_initializer(),
trainable=trainable)
logits = tf.matmul(encoder.get_pooled_output(),
output_weights,
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
probs = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
labels = tf.reshape(next_sentence_labels, [-1])
one_hot_labels = tf.one_hot(labels, depth=2, dtype=tf.float32)
per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs,
axis=-1)
if sample_weight is not None:
per_example_loss = (tf.cast(sample_weight, dtype=tf.float32) *
per_example_loss)
loss = tf.reduce_mean(per_example_loss)
if use_nsp_loss:
scalar_losses.append(loss)
self.losses['NSP_losses'] = per_example_loss
self.probs['NSP_probs'] = probs
self.preds['NSP_preds'] = tf.argmax(probs, axis=-1)
self.total_loss = tf.add_n(scalar_losses)
def __init__(self,
is_training,
input_tensor,
n_wide_features,
wide_features,
label_ids,
label_size=2,
sample_weight=None,
scope='cls/seq_relationship',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
hidden_size = input_tensor.shape.as_list()[-1]
feature_size = wide_features.shape.as_list()[-1]
with tf.variable_scope('wide'):
feature_embeddings = tf.get_variable(
name='feature_embeddings',
shape=[feature_size + 1, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
wide_output = tf.gather(feature_embeddings,
wide_features) # [B, N, H]
with tf.variable_scope('wide_and_deep'):
deep_output = tf.expand_dims(input_tensor, -1) # [B, H, 1]
attention_scores = tf.matmul(wide_output, deep_output) # [B, N, 1]
attention_scores = tf.transpose(attention_scores,
[0, 2, 1]) # [B, 1, N]
attention_scores = tf.multiply(attention_scores,
1.0 / math.sqrt(hidden_size))
feature_mask = tf.cast(
tf.sequence_mask(n_wide_features, feature_size),
tf.float32) # [B, N]
feature_mask = tf.expand_dims(feature_mask, 1) # [B, 1, N]
attention_scores += (1.0 - feature_mask) * -10000.0
attention_matrix = tf.nn.softmax(attention_scores, axis=-1)
attention_output = tf.matmul(attention_matrix,
wide_output) # [B, 1, H]
attention_output = attention_output[:, 0, :] # [B, H]
# attention_output = util.dropout(
# attention_output, hidden_dropout_prob)
input_tensor = util.layer_norm(attention_output + input_tensor,
trainable=trainable)
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
self.preds['preds'] = tf.argmax(logits, axis=-1)
self.probs['probs'] = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(label_ids,
depth=label_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs,
axis=-1)
if sample_weight is not None:
per_example_loss = tf.cast(sample_weight,
dtype=tf.float32) * per_example_loss
thresh = kwargs.get('tsa_thresh')
if thresh is not None:
assert isinstance(
thresh,
float), ('`tsa_thresh` must be a float between 0 and 1.')
uncertainty = tf.reduce_sum(self.probs['probs'] *
tf.log(self.probs['probs']),
axis=-1)
uncertainty /= tf.log(1 / label_size)
per_example_loss = tf.cast(
tf.greater(uncertainty, thresh), dtype=tf.float32) * \
per_example_loss
self.losses['losses'] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
def __init__(self,
is_training,
input_tensor,
is_supervised,
is_expanded,
label_ids,
label_size=2,
sample_weight=None,
scope='cls/seq_relationship',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
global_step=None,
num_train_steps=None,
uda_softmax_temp=-1,
uda_confidence_thresh=-1,
tsa_schedule='linear',
**kwargs):
super().__init__(**kwargs)
is_supervised = tf.cast(is_supervised, tf.float32)
is_expanded = tf.cast(is_expanded, tf.float32)
hidden_size = input_tensor.shape.as_list()[-1]
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
log_probs = tf.nn.log_softmax(logits, axis=-1)
with tf.variable_scope('sup_loss'):
# reshape
sup_ori_log_probs = tf.boolean_mask(log_probs,
mask=(1.0 - is_expanded),
axis=0)
sup_log_probs = tf.boolean_mask(sup_ori_log_probs,
mask=is_supervised,
axis=0)
sup_label_ids = tf.boolean_mask(label_ids,
mask=is_supervised,
axis=0)
self.preds['preds'] = tf.argmax(sup_ori_log_probs, axis=-1)
one_hot_labels = tf.one_hot(sup_label_ids,
depth=label_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(
one_hot_labels * sup_log_probs, axis=-1)
loss_mask = tf.ones_like(per_example_loss, dtype=tf.float32)
correct_label_probs = tf.reduce_sum(one_hot_labels *
tf.exp(sup_log_probs),
axis=-1)
if is_training and tsa_schedule:
tsa_start = 1.0 / label_size
tsa_threshold = get_tsa_threshold(tsa_schedule,
global_step,
num_train_steps,
tsa_start,
end=1)
larger_than_threshold = tf.greater(correct_label_probs,
tsa_threshold)
loss_mask = loss_mask * (
1 - tf.cast(larger_than_threshold, tf.float32))
loss_mask = tf.stop_gradient(loss_mask)
per_example_loss = per_example_loss * loss_mask
if sample_weight is not None:
sup_sample_weight = tf.boolean_mask(sample_weight,
mask=is_supervised,
axis=0)
per_example_loss *= tf.cast(sup_sample_weight,
dtype=tf.float32)
sup_loss = (tf.reduce_sum(per_example_loss) /
tf.maximum(tf.reduce_sum(loss_mask), 1))
self.losses['supervised'] = per_example_loss
with tf.variable_scope('unsup_loss'):
# reshape
ori_log_probs = tf.boolean_mask(sup_ori_log_probs,
mask=(1.0 - is_supervised),
axis=0)
aug_log_probs = tf.boolean_mask(log_probs,
mask=is_expanded,
axis=0)
sup_ori_logits = tf.boolean_mask(logits,
mask=(1.0 - is_expanded),
axis=0)
ori_logits = tf.boolean_mask(sup_ori_logits,
mask=(1.0 - is_supervised),
axis=0)
unsup_loss_mask = 1
if uda_softmax_temp != -1:
tgt_ori_log_probs = tf.nn.log_softmax(ori_logits /
uda_softmax_temp,
axis=-1)
tgt_ori_log_probs = tf.stop_gradient(tgt_ori_log_probs)
else:
tgt_ori_log_probs = tf.stop_gradient(ori_log_probs)
if uda_confidence_thresh != -1:
largest_prob = tf.reduce_max(tf.exp(ori_log_probs),
axis=-1)
unsup_loss_mask = tf.cast(
tf.greater(largest_prob, uda_confidence_thresh),
tf.float32)
unsup_loss_mask = tf.stop_gradient(unsup_loss_mask)
per_example_loss = kl_for_log_probs(
tgt_ori_log_probs, aug_log_probs) * unsup_loss_mask
if sample_weight is not None:
unsup_sample_weight = tf.boolean_mask(sample_weight,
mask=(1.0 -
is_supervised),
axis=0)
per_example_loss *= tf.cast(unsup_sample_weight,
dtype=tf.float32)
unsup_loss = tf.reduce_mean(per_example_loss)
self.losses['unsupervised'] = per_example_loss
self.total_loss = sup_loss + unsup_loss
def layer_norm(input_tensor,
center=True,
scale=True,
activation_fn=None,
variables_collections=None,
outputs_collections=None,
begin_norm_axis=-1,
begin_params_axis=-1,
trainable=True):
''' Runs layer normalization on the last dimension of the tensor.
Args:
input_tensor: A tensor having rank `R`. The normalization is performed
over axes `begin_norm_axis ... R - 1` and centering and scaling
parameters are calculated over `begin_params_axis ... R - 1`.
center: If True, add offset of `beta` to normalized tensor. If False,
`beta` is ignored.
scale: If True, multiply by `gamma`. If False, `gamma` is not used.
When the next layer is linear (also e.g. `nn.relu`), this can be
disabled since the scaling can be done by the next layer.
activation_fn: Activation function, default set to None to skip it and
maintain a linear activation.
variables_collections: Optional collections for the variables.
outputs_collections: Collections to add the outputs.
trainable: If `True` also add variables to the graph collection
`GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
begin_norm_axis: The first normalization dimension: normalization will
be performed along dimensions `begin_norm_axis : rank(input_tensor)`
begin_params_axis: The first parameter (beta, gamma) dimension: scale
and centering parameters will have dimensions
`begin_params_axis : rank(input_tensor)` and will be broadcast with
the normalized inputs accordingly.
scope: Optional scope for `variable_scope`.
Returns:
A `Tensor` representing the output of the operation, having the same
shape and dtype as `input_tensor`.
Raises:
ValueError: If the rank of `input_tensor` is not known at graph build
time, or if `input_tensor.shape[begin_params_axis:]` is not fully
defined at graph build time.
'''
with tf.variable_scope('LayerNorm'):
inputs_shape = input_tensor.shape
inputs_rank = inputs_shape.ndims
if inputs_rank is None:
raise ValueError('Inputs %s has undefined rank.'
% input_tensor.name)
dtype = input_tensor.dtype.base_dtype
if begin_norm_axis < 0:
begin_norm_axis = inputs_rank + begin_norm_axis
if begin_params_axis >= inputs_rank or begin_norm_axis >= inputs_rank:
raise ValueError(
'begin_params_axis (%d) and begin_norm_axis (%d) '
'must be < rank(inputs) (%d)'
% (begin_params_axis, begin_norm_axis, inputs_rank))
params_shape = inputs_shape[begin_params_axis:]
if not params_shape.is_fully_defined():
raise ValueError(
'Inputs %s: shape(inputs)[%s:] is not fully defined: %s' %
(input_tensor.name, begin_params_axis, inputs_shape))
# Allocate parameters for the beta and gamma of the normalization.
beta, gamma = None, None
if center:
beta = tf.get_variable(
'beta', shape=params_shape,
dtype=dtype,
initializer=tf.zeros_initializer(),
trainable=trainable)
if scale:
gamma = tf.get_variable(
'gamma',
shape=params_shape,
dtype=dtype,
initializer=tf.ones_initializer(),
trainable=trainable)
# By default, compute the moments across all the dimensions except the
# one with index 0.
norm_axes = list(range(begin_norm_axis, inputs_rank))
mean, variance = tf.nn.moments(input_tensor, norm_axes, keep_dims=True)
# Compute layer normalization using the batch_normalization function.
# Note that epsilon must be increased for float16 due to the limited
# representable range.
variance_epsilon = 1e-12 if dtype != tf.float16 else 1e-3
outputs = tf.nn.batch_normalization(
input_tensor, mean, variance,
offset=beta, scale=gamma,
variance_epsilon=variance_epsilon)
outputs.set_shape(inputs_shape)
if activation_fn is not None:
outputs = activation_fn(outputs)
return outputs
def __init__(self,
bert_config,
is_training,
sketchy_encoder,
intensive_encoder,
query_mask,
label_ids,
has_answer,
sample_weight=None,
scope='retro_reader',
matching_mechanism='cross-attention',
beta_1=0.5,
beta_2=0.5,
threshold=1.0,
trainable=True,
**kwargs):
super().__init__(**kwargs)
# verifier
with tf.variable_scope(scope):
# sketchy reading module
with tf.variable_scope('sketchy/prediction'):
sketchy_output = sketchy_encoder.get_pooled_output()
hidden_size = sketchy_output.shape.as_list()[-1]
output_weights = tf.get_variable(
'output_weights',
shape=[2, hidden_size],
initializer=util.create_initializer(
bert_config.initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[2],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
sketchy_output, bert_config.hidden_dropout_prob \
if is_training else 0.0)
logits = tf.matmul(
output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(
has_answer, depth=2, dtype=tf.float32)
per_example_loss = - tf.reduce_sum(
one_hot_labels * log_probs, axis=-1)
if sample_weight is not None:
per_example_loss = tf.cast(
sample_weight, dtype=tf.float32) * per_example_loss
self.losses['sketchy_losses'] = per_example_loss
sketchy_loss = tf.reduce_mean(per_example_loss)
score_ext = logits[:, 1] - logits[:, 0]
# intensive reading module
with tf.variable_scope('intensive'):
H = intensive_encoder.get_sequence_output()
H_Q = H * tf.cast(
tf.expand_dims(query_mask, axis=-1), tf.float32)
(batch_size, max_seq_length, hidden_size) = \
util.get_shape_list(H)
# cross-attention
if matching_mechanism == 'cross-attention':
with tf.variable_scope('cross_attention'):
attention_mask = \
self.create_attention_mask_from_input_mask(
query_mask, batch_size, max_seq_length)
(H_prime, _) = self.attention_layer(
from_tensor=H,
to_tensor=H_Q,
attention_mask=attention_mask,
num_attention_heads=\
bert_config.num_attention_heads,
size_per_head=\
hidden_size // bert_config.num_attention_heads,
attention_probs_dropout_prob=\
bert_config.hidden_dropout_prob,
initializer_range=bert_config.initializer_range,
do_return_2d_tensor=False,
batch_size=batch_size,
from_max_seq_length=max_seq_length,
to_max_seq_length=max_seq_length,
trainable=trainable)
# matching-attention
elif matching_mechanism == 'matching-attention':
with tf.variable_scope('matching_attention'):
output_weights = tf.get_variable(
'output_weights',
shape=[hidden_size, hidden_size],
initializer=util.create_initializer(
bert_config.initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[hidden_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
trans = tf.matmul(
H_Q, tf.tile(
tf.expand_dims(output_weights, axis=0),
[batch_size, 1, 1]),
transpose_b=True)
trans = tf.nn.bias_add(trans, output_bias)
M = tf.nn.softmax(
tf.matmul(H, trans, transpose_b=True), axis=-1)
H_prime = tf.matmul(M, H_Q)
with tf.variable_scope('prediction'):
output_weights = tf.get_variable(
'output_weights',
shape=[2, hidden_size],
initializer=util.create_initializer(
bert_config.initializer_range),
trainable=trainable)
output_bias = tf.get_variable(
'output_bias',
shape=[2],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
H_prime, bert_config.hidden_dropout_prob \
if is_training else 0.0)
output_layer = tf.reshape(
output_layer,
[batch_size * max_seq_length, hidden_size])
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(
logits, [batch_size, max_seq_length, 2])
logits = tf.transpose(logits, [0, 2, 1])
probs = tf.nn.softmax(logits, axis=-1, name='probs')
self.probs['mrc_probs'] = probs
self.preds['mrc_preds'] = tf.argmax(logits, axis=-1)
start_one_hot_labels = tf.one_hot(
label_ids[:, 0], depth=max_seq_length,
dtype=tf.float32)
end_one_hot_labels = tf.one_hot(
label_ids[:, 1], depth=max_seq_length,
dtype=tf.float32)
start_log_probs = tf.nn.log_softmax(logits[:, 0, :], axis=-1)
end_log_probs = tf.nn.log_softmax(logits[:, 1, :], axis=-1)
per_example_loss = (
- 0.5 * tf.reduce_sum(
start_one_hot_labels * start_log_probs, axis=-1)
- 0.5 * tf.reduce_sum(
end_one_hot_labels * end_log_probs, axis=-1))
if sample_weight is not None:
per_example_loss *= sample_weight
intensive_loss = tf.reduce_mean(per_example_loss)
self.losses['intensive_losses'] = per_example_loss
score_has = tf.norm(
probs[:, 0, 1:] + probs[:, 1, 1:], np.inf, axis=-1)
score_null = probs[:, 0, 0] + probs[:, 1, 0]
score_diff = score_has - score_null
# rear verification
v = beta_1 * score_diff + beta_2 * score_ext
self.preds['verifier_preds'] = \
tf.cast(tf.greater(v, threshold), tf.int32)
self.probs['verifier_probs'] = v
self.total_loss = sketchy_loss + intensive_loss
def __init__(self,
bert_config,
is_training,
input_ids,
input_mask,
segment_ids,
sample_weight=None,
scope='bert',
dtype=tf.float32,
drop_pooler=False,
cls_model='self-attention',
label_size=2,
speed=0.1,
ignore_cls='0',
**kwargs):
super(FastBERTCLSDistillor, self).__init__()
if not ignore_cls:
ignore_cls = []
if isinstance(ignore_cls, str):
ignore_cls = ignore_cls.replace(' ', '').split(',')
ignore_cls = list(map(int, ignore_cls))
elif isinstance(ignore_cls, list):
ignore_cls = list(map(int, ignore_cls))
else:
raise ValueError(
'`ignore_cls` should be a list of child-classifier ids or '
'a string seperated with commas.')
if not speed:
raise ValueError(
'`speed` should be a float number between `0` and `1`.')
bert_config = copy.deepcopy(bert_config)
bert_config.hidden_dropout_prob = 0.0
bert_config.attention_probs_dropout_prob = 0.0
input_shape = util.get_shape_list(input_ids, expected_rank=2)
batch_size = input_shape[0]
max_seq_length = input_shape[1]
with tf.variable_scope(scope):
with tf.variable_scope('embeddings'):
(self.embedding_output, self.embedding_table) = \
self.embedding_lookup(
input_ids=input_ids,
vocab_size=bert_config.vocab_size,
batch_size=batch_size,
max_seq_length=max_seq_length,
embedding_size=bert_config.hidden_size,
initializer_range=bert_config.initializer_range,
word_embedding_name='word_embeddings',
dtype=dtype,
trainable=False,
tilda_embeddings=None)
# Add positional embeddings and token type embeddings
# layer normalize and perform dropout.
self.embedding_output = self.embedding_postprocessor(
input_tensor=self.embedding_output,
batch_size=batch_size,
max_seq_length=max_seq_length,
hidden_size=bert_config.hidden_size,
use_token_type=True,
segment_ids=segment_ids,
token_type_vocab_size=bert_config.type_vocab_size,
token_type_embedding_name='token_type_embeddings',
use_position_embeddings=True,
position_embedding_name='position_embeddings',
initializer_range=bert_config.initializer_range,
max_position_embeddings=\
bert_config.max_position_embeddings,
dropout_prob=bert_config.hidden_dropout_prob,
dtype=dtype,
trainable=False)
with tf.variable_scope('encoder'):
attention_mask = self.create_attention_mask_from_input_mask(
input_mask, batch_size, max_seq_length, dtype=dtype)
# stacked transformers
(self.all_encoder_layers, self.all_cls_layers) = \
self.dynamic_transformer_model(
is_training,
input_tensor=self.embedding_output,
input_mask=input_mask,
batch_size=batch_size,
max_seq_length=max_seq_length,
label_size=label_size,
attention_mask=attention_mask,
hidden_size=bert_config.hidden_size,
num_hidden_layers=bert_config.num_hidden_layers,
num_attention_heads=bert_config.num_attention_heads,
intermediate_size=bert_config.intermediate_size,
intermediate_act_fn=util.get_activation(
bert_config.hidden_act),
hidden_dropout_prob=bert_config.hidden_dropout_prob,
attention_probs_dropout_prob=\
bert_config.attention_probs_dropout_prob,
initializer_range=bert_config.initializer_range,
dtype=dtype,
cls_model=cls_model,
speed=speed,
ignore_cls=ignore_cls)
self.sequence_output = self.all_encoder_layers[-1]
with tf.variable_scope('pooler'):
first_token_tensor = self.sequence_output[:, 0, :]
# trick: ignore the fully connected layer
if drop_pooler:
self.pooled_output = first_token_tensor
else:
self.pooled_output = tf.layers.dense(
first_token_tensor,
bert_config.hidden_size,
activation=tf.tanh,
kernel_initializer=util.create_initializer(
bert_config.initializer_range),
trainable=False)
# teacher classifier
if bert_config.num_hidden_layers not in ignore_cls:
with tf.variable_scope('cls/seq_relationship'):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, bert_config.hidden_size],
initializer=util.create_initializer(
bert_config.initializer_range),
trainable=False)
output_bias = tf.get_variable(
'output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=False)
logits = tf.matmul(self.pooled_output,
output_weights,
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
probs = tf.nn.softmax(logits, axis=-1)
# distillation
if is_training:
losses = []
for cls_probs in self.all_cls_layers.values():
# KL-Divergence
per_example_loss = tf.reduce_sum(
cls_probs * (tf.log(cls_probs) - tf.log(probs)), axis=-1)
if sample_weight is not None:
per_example_loss *= tf.cast(sample_weight,
dtype=tf.float32)
loss = tf.reduce_mean(per_example_loss)
losses.append(loss)
distill_loss = tf.add_n(losses)
self.total_loss = distill_loss
self.losses['losses'] = distill_loss
else:
if bert_config.num_hidden_layers not in ignore_cls:
self.all_cls_layers[bert_config.num_hidden_layers] = probs
self.probs['probs'] = tf.concat(list(self.all_cls_layers.values()),
axis=0,
name='probs')
def __init__(self,
vocab_size,
is_training,
input_ids,
input_mask,
segment_ids,
sample_weight=None,
reduced_size=64,
topic_size=1024,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
bias=0,
scope='vae',
trainable=True,
**kwargs):
super().__init__()
# freeze parameters
config = Config(vocab_size,
hidden_size=hidden_size,
num_hidden_layers=num_hidden_layers,
num_attention_heads=num_attention_heads)
if not is_training:
config.hidden_dropout_prob = 0.0
config.attention_probs_dropout_prob = 0.0
input_shape = util.get_shape_list(input_ids, expected_rank=2)
batch_size = input_shape[0]
seq_length = input_shape[1]
# Tilda embeddings for SMART algorithm
tilda_embeddings = None
use_tilda_embedding = kwargs.get('use_tilda_embedding')
if use_tilda_embedding:
with tf.variable_scope('', reuse=True):
tilda_embeddings = tf.get_variable('tilda_embeddings')
with tf.variable_scope(scope):
with tf.variable_scope('embeddings'):
(self.embedding_output, self.embedding_table) = \
self.embedding_lookup(
input_ids=input_ids,
vocab_size=config.vocab_size,
batch_size=batch_size,
max_seq_length=seq_length,
embedding_size=config.hidden_size,
initializer_range=config.initializer_range,
word_embedding_name='word_embeddings',
tilda_embeddings=tilda_embeddings,
trainable=trainable)
self.embedding_output = self.embedding_postprocessor(
input_tensor=self.embedding_output,
batch_size=batch_size,
max_seq_length=seq_length,
hidden_size=config.hidden_size,
use_token_type=True,
segment_ids=segment_ids,
token_type_vocab_size=config.type_vocab_size,
token_type_embedding_name='token_type_embeddings',
use_position_embeddings=True,
position_embedding_name='position_embeddings',
initializer_range=config.initializer_range,
max_position_embeddings=config.max_position_embeddings,
dropout_prob=config.hidden_dropout_prob,
trainable=trainable)
with tf.variable_scope('encoder'):
# stacked transformer
attention_mask = self.create_attention_mask_from_input_mask(
input_mask, batch_size, seq_length)
self.all_encoder_layers = self.transformer_model(
input_tensor=self.embedding_output,
batch_size=batch_size,
max_seq_length=seq_length,
attention_mask=attention_mask,
hidden_size=config.hidden_size,
num_hidden_layers=config.num_hidden_layers,
num_attention_heads=config.num_attention_heads,
intermediate_size=config.intermediate_size,
intermediate_act_fn=util.get_activation(config.hidden_act),
hidden_dropout_prob=config.hidden_dropout_prob,
attention_probs_dropout_prob=\
config.attention_probs_dropout_prob,
initializer_range=config.initializer_range,
trainable=trainable)
# projection
with tf.variable_scope('projection'):
transformer_output = tf.layers.dense(
self.all_encoder_layers[-1],
reduced_size,
activation=util.gelu,
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
trainable=trainable)
transformer_output = tf.reshape(transformer_output,
[batch_size, -1])
input_length = tf.reduce_sum(input_mask, axis=-1)
input_length = tf.cast(input_length, tf.float32)
input_length_1d = tf.reshape(input_length, [batch_size])
input_length_2d = tf.reshape(input_length, [batch_size, 1])
broadcast_mask = tf.sequence_mask(
tf.multiply(input_length_1d, reduced_size),
seq_length * reduced_size,
dtype=tf.float32)
broadcast_mask = tf.multiply(broadcast_mask,
seq_length / input_length_2d)
transformer_output *= broadcast_mask
# latent space
miu = tf.layers.dense(
transformer_output,
topic_size,
activation='tanh',
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
name='miu',
trainable=trainable)
sigma = tf.layers.dense(
transformer_output,
topic_size,
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
name='sigma',
trainable=trainable)
self.probs['miu'] = miu
self.probs['sigma'] = sigma
with tf.variable_scope('decoder'):
with tf.variable_scope('projection'):
# reparametarization
if is_training:
noise = tf.random_normal([batch_size, topic_size])
else:
noise = tf.random_uniform([batch_size, topic_size],
minval=-bias,
maxval=bias)
decoder_input = miu + tf.exp(sigma) * noise
# projection
decoder_input = tf.layers.dense(
decoder_input,
seq_length * reduced_size,
activation=util.gelu,
kernel_initializer=tf.truncated_normal_initializer(
stddev=config.initializer_range),
trainable=trainable)
intermediate_input = tf.reshape(
decoder_input, [-1, seq_length, reduced_size])
intermediate_input = util.layer_norm(intermediate_input,
trainable=trainable)
intermediate_input = util.dropout(
intermediate_input, config.hidden_dropout_prob)
# MLP
with tf.variable_scope('intermediate'):
intermediate_output = tf.layers.dense(
intermediate_input,
4 * reduced_size,
activation=util.gelu,
kernel_initializer=util.create_initializer(
config.initializer_range),
trainable=trainable)
with tf.variable_scope('output'):
decoder_output = tf.layers.dense(
intermediate_output,
config.hidden_size,
kernel_initializer=util.create_initializer(
config.initializer_range),
trainable=trainable)
decoder_output = util.layer_norm(decoder_output,
trainable=trainable)
decoder_output = util.dropout(decoder_output,
config.hidden_dropout_prob)
self.all_decoder_layers = [intermediate_output, decoder_output]
self.all_decoder_layers = [decoder_output]
# reconstruction
with tf.variable_scope('cls/predictions'):
with tf.variable_scope('transform'):
input_tensor = tf.layers.dense(
decoder_output,
units=config.hidden_size,
activation=util.get_activation(config.hidden_act),
kernel_initializer=util.create_initializer(
config.initializer_range),
trainable=trainable)
input_tensor = util.layer_norm(input_tensor,
trainable=trainable)
output_weights = self.embedding_table
output_bias = tf.get_variable('output_bias',
shape=[config.vocab_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
flatten_input_tensor = tf.reshape(input_tensor,
[-1, config.hidden_size])
logits = tf.matmul(flatten_input_tensor,
output_weights,
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
logits = tf.reshape(logits,
[batch_size, seq_length, config.vocab_size])
probs = tf.nn.softmax(logits, axis=-1, name='probs')
lm_log_probs = tf.nn.log_softmax(logits, axis=-1)
self.preds['preds'] = tf.argmax(probs, axis=-1)
one_hot_labels = tf.one_hot(input_ids,
depth=config.vocab_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(lm_log_probs * one_hot_labels,
axis=[-1])
if sample_weight is not None:
per_example_loss *= tf.expand_dims(sample_weight, axis=-1)
self.total_loss = (tf.reduce_mean(per_example_loss) +
tf.reduce_mean(tf.square(miu)) +
tf.reduce_mean(tf.exp(sigma) - sigma - 1))
self.losses['losses'] = per_example_loss
def __init__(self,
bert_config,
is_training,
input_tensor,
input_mask,
sem_features,
label_ids,
max_seq_length,
feature_size,
label_size=2,
sample_weight=None,
scope='cls/seq_relationship',
hidden_dropout_prob=0.1,
initializer_range=0.02,
trainable=True,
**kwargs):
super().__init__(**kwargs)
input_shape = util.get_shape_list(input_tensor)
batch_size = input_shape[0]
hidden_size = input_shape[-1]
with tf.variable_scope('sem'):
feature_embeddings = tf.get_variable(
name='feature_embeddings',
shape=[feature_size + 3,
hidden_size], # for [PAD], [CLS], [SEP]
initializer=util.create_initializer(initializer_range),
trainable=trainable)
sem_output = tf.gather(feature_embeddings,
sem_features) # [B, N, H]
attention_heads = []
with tf.variable_scope('self'):
attention_mask = BERTEncoder.create_attention_mask_from_input_mask(
input_mask, batch_size, max_seq_length)
(attention_head, _) = BERTEncoder.attention_layer(
from_tensor=sem_output,
to_tensor=sem_output,
attention_mask=attention_mask,
num_attention_heads=bert_config.num_attention_heads,
size_per_head=(hidden_size //
bert_config.num_attention_heads),
attention_probs_dropout_prob=hidden_dropout_prob
if is_training else 0.0,
initializer_range=initializer_range,
do_return_2d_tensor=False,
batch_size=batch_size,
from_max_seq_length=max_seq_length,
to_max_seq_length=max_seq_length,
trainable=trainable)
attention_heads.append(attention_head)
if len(attention_heads) == 1:
attention_output = attention_heads[0]
else:
attention_output = tf.concat(attention_heads, axis=-1)
attention_output = attention_output[:, 0, :] # [B, H]
input_tensor = util.layer_norm(attention_output + input_tensor,
trainable=trainable)
with tf.variable_scope(scope):
output_weights = tf.get_variable(
'output_weights',
shape=[label_size, hidden_size],
initializer=util.create_initializer(initializer_range),
trainable=trainable)
output_bias = tf.get_variable('output_bias',
shape=[label_size],
initializer=tf.zeros_initializer(),
trainable=trainable)
output_layer = util.dropout(
input_tensor, hidden_dropout_prob if is_training else 0.0)
logits = tf.matmul(output_layer, output_weights, transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
self.preds['preds'] = tf.argmax(logits, axis=-1)
self.probs['probs'] = tf.nn.softmax(logits, axis=-1, name='probs')
log_probs = tf.nn.log_softmax(logits, axis=-1)
one_hot_labels = tf.one_hot(label_ids,
depth=label_size,
dtype=tf.float32)
per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs,
axis=-1)
if sample_weight is not None:
per_example_loss = tf.cast(sample_weight,
dtype=tf.float32) * per_example_loss
thresh = kwargs.get('tsa_thresh')
if thresh is not None:
assert isinstance(
thresh,
float), ('`tsa_thresh` must be a float between 0 and 1.')
uncertainty = tf.reduce_sum(self.probs['probs'] *
tf.log(self.probs['probs']),
axis=-1)
uncertainty /= tf.log(1 / label_size)
per_example_loss = tf.cast(
tf.greater(uncertainty, thresh), dtype=tf.float32) * \
per_example_loss
self.losses['losses'] = per_example_loss
self.total_loss = tf.reduce_mean(per_example_loss)
