def get_classification_outputs(FLAGS, features, is_training):
"""Loss for downstream classification tasks."""
input_ids = features["input_ids"]
seg_id = features["segment_ids"]
input_mask_int = tf.cast(tf.cast(input_ids, tf.bool), tf.int32)
input_mask = 1 - tf.cast(input_mask_int, tf.float32)
num_choices = FLAGS.num_choices
batch_size = tf.shape(features["input_ids"])[0]
def _transform_features(feature):
out = tf.reshape(feature, [batch_size, num_choices, -1])
out = tf.transpose(out, [2, 0, 1])
out = tf.reshape(out, [-1, batch_size * num_choices])
return out
if num_choices:
input_ids = _transform_features(input_ids)
seg_id = _transform_features(seg_id)
input_mask = _transform_features(input_mask)
else:
input_ids = tf.transpose(input_ids, [1, 0])
seg_id = tf.transpose(seg_id, [1, 0])
input_mask = tf.transpose(input_mask, [1, 0])
xlnet_config = xlnet.XLNetConfig(json_path=FLAGS.model_config_path)
run_config = xlnet.create_run_config(is_training, True, FLAGS)
xlnet_model = xlnet.XLNetModel(
xlnet_config=xlnet_config,
run_config=run_config,
input_ids=input_ids,
seg_ids=seg_id,
input_mask=input_mask)
summary = xlnet_model.get_pooled_out(FLAGS.summary_type,
FLAGS.use_summ_proj)
initializer = xlnet_model.get_initializer()
return_dict = {}
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
with tf.variable_scope("answer_class"):
# race has 4 classes,
# boolq has 2 classes
if num_choices:
num_classes = 1
else:
num_classes = FLAGS.num_classes
cls_logits = tf.layers.dense(summary, num_classes,
kernel_initializer=initializer,
name="cls")
if num_choices:
cls_logits = tf.reshape(cls_logits, [batch_size, num_choices])
cls_log_probs = tf.nn.log_softmax(cls_logits, -1)
if is_training:
return_dict["cls_log_probs"] = cls_log_probs
return_dict["cls_logits"] = cls_logits
return return_dict
def multihead_attn(q,
k,
v,
attn_mask,
d_model,
n_head,
d_head,
dropout,
dropatt,
is_training,
kernel_initializer,
residual=True,
scope='abs_attn',
reuse=None):
"""Standard multi-head attention with absolute positional embedding."""
scale = 1 / (d_head**0.5)
with tf.variable_scope(scope, reuse=reuse):
# attention heads
q_head = head_projection(q, d_model, n_head, d_head,
kernel_initializer, 'q')
k_head = head_projection(k, d_model, n_head, d_head,
kernel_initializer, 'k')
v_head = head_projection(v, d_model, n_head, d_head,
kernel_initializer, 'v')
# attention vector
attn_vec = abs_attn_core(q_head, k_head, v_head, attn_mask, dropatt,
is_training, scale)
# post processing
output = post_attention(v, attn_vec, d_model, n_head, d_head, dropout,
is_training, kernel_initializer, residual)
return output
def classification_loss(hidden,
labels,
n_class,
initializer,
scope,
reuse=None,
return_logits=False):
"""
Different classification tasks should use different scope names to ensure
different dense layers (parameters) are used to produce the logits.
An exception will be in transfer learning, where one hopes to transfer
the classification weights.
"""
with tf.variable_scope(scope, reuse=reuse):
logits = tf.layers.dense(hidden,
n_class,
kernel_initializer=initializer,
name='logit')
one_hot_target = tf.one_hot(labels, n_class, dtype=hidden.dtype)
loss = -tf.reduce_sum(tf.nn.log_softmax(logits) * one_hot_target, -1)
if return_logits:
return loss, logits
return loss
def get_pooled_out(self, summary_type, use_summ_proj=True):
"""
Args:
summary_type: str, "last", "first", "mean", or "attn". The method
to pool the input to get a vector representation.
use_summ_proj: bool, whether to use a linear projection during pooling.
Returns:
float32 Tensor in shape [bsz, d_model], the pooled representation.
"""
xlnet_config = self.xlnet_config
run_config = self.run_config
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
summary = modeling.summarize_sequence(
summary_type=summary_type,
hidden=self.output,
d_model=xlnet_config.d_model,
n_head=xlnet_config.n_head,
d_head=xlnet_config.d_head,
dropout=run_config.dropout,
dropatt=run_config.dropatt,
is_training=run_config.is_training,
input_mask=self.input_mask,
initializer=self.initializer,
use_proj=use_summ_proj)
return summary
def rel_multihead_attn(h,
r,
r_w_bias,
r_r_bias,
seg_mat,
r_s_bias,
seg_embed,
attn_mask,
mems,
d_model,
n_head,
d_head,
dropout,
dropatt,
is_training,
kernel_initializer,
scope='rel_attn',
reuse=None):
"""Multi-head attention with relative positional encoding."""
scale = 1 / (d_head**0.5)
with tf.variable_scope(scope, reuse=reuse):
if mems is not None and mems.shape.ndims > 1:
cat = tf.concat([mems, h], 0)
else:
cat = h
# content heads
q_head_h = head_projection(h, d_model, n_head, d_head,
kernel_initializer, 'q')
k_head_h = head_projection(cat, d_model, n_head, d_head,
kernel_initializer, 'k')
v_head_h = head_projection(cat, d_model, n_head, d_head,
kernel_initializer, 'v')
# positional heads
k_head_r = head_projection(r, d_model, n_head, d_head,
kernel_initializer, 'r')
# core attention ops
attn_vec = rel_attn_core(q_head_h, k_head_h, v_head_h, k_head_r,
seg_embed, seg_mat, r_w_bias, r_r_bias,
r_s_bias, attn_mask, dropatt, is_training,
scale)
# post processing
output = post_attention(h, attn_vec, d_model, n_head, d_head, dropout,
is_training, kernel_initializer)
return output
def clean_ckpt(_):
input_ckpt = FLAGS.clean_input_ckpt
output_model_dir = FLAGS.clean_output_model_dir
tf.reset_default_graph()
var_list = tf.contrib.framework.list_variables(input_ckpt)
var_values, var_dtypes = {}, {}
for (name, shape) in var_list:
if not name.startswith("global_step") and "adam" not in name.lower():
var_values[name] = None
logger.info("Include {}".format(name))
else:
logger.info("Exclude {}".format(name))
logger.info("Loading from {}".format(input_ckpt))
reader = tf.contrib.framework.load_checkpoint(input_ckpt)
for name in var_values:
tensor = reader.get_tensor(name)
var_dtypes[name] = tensor.dtype
var_values[name] = tensor
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
tf_vars = [
tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
for v in var_values
]
placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
global_step = tf.Variable(0,
name="global_step",
trainable=False,
dtype=tf.int64)
saver = tf.train.Saver(tf.all_variables())
if not tf.gfile.Exists(output_model_dir):
tf.gfile.MakeDirs(output_model_dir)
# Build a model consisting only of variables, set them to the average values.
with tf.Session() as sess:
sess.run(tf.initialize_all_variables())
for p, assign_op, (name, value) in zip(placeholders, assign_ops,
six.iteritems(var_values)):
sess.run(assign_op, {p: value})
# Use the built saver to save the averaged checkpoint.
saver.save(sess,
join(output_model_dir, "model.ckpt"),
global_step=global_step)
def get_decomposed_classification_outputs(FLAGS, features, is_training):
seq1_ids = features["seq1_ids"]
seq2_ids = features["seq2_ids"]
seq_len = FLAGS.max_seq_length
first_seq_len = FLAGS.max_first_length + 2
second_seq_len = seq_len - first_seq_len
seq1_attn_mask = get_attention_mask(seq1_ids, first_seq_len)
seq2_attn_mask = get_attention_mask(seq2_ids, second_seq_len)
seq_attn_mask = get_attention_mask(tf.concat([seq2_ids, seq1_ids], axis=0),
seq_len)
xlnet_config = xlnet.XLNetConfig(json_path=FLAGS.model_config_path)
run_config = xlnet.create_run_config(is_training, True, FLAGS)
initializer = xlnet._get_initializer(run_config)
tfm_args = dict(
n_token=xlnet_config.n_token,
initializer=initializer,
attn_type="bi",
n_layer=xlnet_config.n_layer,
d_model=xlnet_config.d_model,
n_head=xlnet_config.n_head,
d_head=xlnet_config.d_head,
d_inner=xlnet_config.d_inner,
ff_activation=xlnet_config.ff_activation,
untie_r=xlnet_config.untie_r,
is_training=run_config.is_training,
use_bfloat16=run_config.use_bfloat16,
use_tpu=run_config.use_tpu,
dropout=run_config.dropout,
dropatt=run_config.dropatt,
# mem_len=run_config.mem_len,
# reuse_len=run_config.reuse_len,
# bi_data=run_config.bi_data,
clamp_len=run_config.clamp_len,
# same_length=run_config.same_length,
ctx_ids=seq2_ids,
q_ids=seq1_ids,
q_seq_len=first_seq_len,
ctx_seq_len=second_seq_len,
sep_layer=FLAGS.sep_layer,
q_attn_mask=seq1_attn_mask,
c_attn_mask=seq2_attn_mask,
qc_attn_mask=seq_attn_mask,
)
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
upper_outputs = transformer_xl_decomposed(**tfm_args)
def avg_checkpoints(model_dir, output_model_dir, last_k):
tf.reset_default_graph()
checkpoint_state = tf.train.get_checkpoint_state(model_dir)
checkpoints = checkpoint_state.all_model_checkpoint_paths[-last_k:]
var_list = tf.contrib.framework.list_variables(checkpoints[0])
var_values, var_dtypes = {}, {}
for (name, shape) in var_list:
if not name.startswith("global_step"):
var_values[name] = np.zeros(shape)
for checkpoint in checkpoints:
reader = tf.contrib.framework.load_checkpoint(checkpoint)
for name in var_values:
tensor = reader.get_tensor(name)
var_dtypes[name] = tensor.dtype
var_values[name] += tensor
logger.info("Read from checkpoint %s", checkpoint)
for name in var_values: # Average.
var_values[name] /= len(checkpoints)
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
tf_vars = [
tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
for v in var_values
]
placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
global_step = tf.Variable(0,
name="global_step",
trainable=False,
dtype=tf.int64)
saver = tf.train.Saver(tf.all_variables())
# Build a model consisting only of variables, set them to the average values.
with tf.Session() as sess:
sess.run(tf.initialize_all_variables())
for p, assign_op, (name, value) in zip(placeholders, assign_ops,
six.iteritems(var_values)):
sess.run(assign_op, {p: value})
# Use the built saver to save the averaged checkpoint.
saver.save(sess,
join(output_model_dir, "model.ckpt"),
global_step=global_step)
def positionwise_ffn(inp,
d_model,
d_inner,
dropout,
kernel_initializer,
activation_type='relu',
scope='ff',
is_training=True,
reuse=None):
"""Position-wise Feed-forward Network."""
if activation_type == 'relu':
activation = tf.nn.relu
elif activation_type == 'gelu':
activation = gelu
else:
raise ValueError(
'Unsupported activation type {}'.format(activation_type))
output = inp
with tf.variable_scope(scope, reuse=reuse):
output = tf.layers.dense(output,
d_inner,
activation=activation,
kernel_initializer=kernel_initializer,
name='layer_1')
output = tf.layers.dropout(output,
dropout,
training=is_training,
name='drop_1')
output = tf.layers.dense(output,
d_model,
kernel_initializer=kernel_initializer,
name='layer_2')
output = tf.layers.dropout(output,
dropout,
training=is_training,
name='drop_2')
output = tf.contrib.layers.layer_norm(output + inp,
begin_norm_axis=-1,
scope='LayerNorm')
return output
def get_race_loss(FLAGS, features, is_training):
"""Loss for downstream multi-choice QA tasks such as RACE."""
bsz_per_core = tf.shape(features["input_ids"])[0]
def _transform_features(feature):
out = tf.reshape(feature, [bsz_per_core, 4, -1])
out = tf.transpose(out, [2, 0, 1])
out = tf.reshape(out, [-1, bsz_per_core * 4])
return out
inp = _transform_features(features["input_ids"])
seg_id = _transform_features(features["segment_ids"])
inp_mask = _transform_features(features["input_mask"])
label = tf.reshape(features["label_ids"], [bsz_per_core])
xlnet_config = xlnet.XLNetConfig(json_path=FLAGS.model_config_path)
run_config = xlnet.create_run_config(is_training, True, FLAGS)
xlnet_model = xlnet.XLNetModel(
xlnet_config=xlnet_config,
run_config=run_config,
input_ids=inp,
seg_ids=seg_id,
input_mask=inp_mask)
summary = xlnet_model.get_pooled_out(FLAGS.summary_type,
FLAGS.use_summ_proj)
with tf.variable_scope("logits"):
logits = tf.layers.dense(summary, 1,
kernel_initializer=xlnet_model.get_initializer())
logits = tf.reshape(logits, [bsz_per_core, 4])
one_hot_target = tf.one_hot(label, 4)
per_example_loss = -tf.reduce_sum(
tf.nn.log_softmax(logits) * one_hot_target, -1)
total_loss = tf.reduce_mean(per_example_loss)
return total_loss, per_example_loss, logits
def embedding_lookup(x,
n_token,
d_embed,
initializer,
use_tpu=True,
scope='embedding',
reuse=None,
dtype=tf.float32):
"""TPU and GPU embedding_lookup function."""
with tf.variable_scope(scope, reuse=reuse):
lookup_table = tf.get_variable('lookup_table', [n_token, d_embed],
dtype=dtype,
initializer=initializer)
if use_tpu:
one_hot_idx = tf.one_hot(x, n_token, dtype=dtype)
if one_hot_idx.shape.ndims == 2:
return tf.einsum('in,nd->id', one_hot_idx,
lookup_table), lookup_table
else:
return tf.einsum('ibn,nd->ibd', one_hot_idx,
lookup_table), lookup_table
else:
return tf.nn.embedding_lookup(lookup_table, x), lookup_table
def model_fn(features, labels, mode, params):
# ### Training or Evaluation
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
# ### Get loss from inputs
sep_layer = FLAGS.sep_layer
if FLAGS.supervise:
FLAGS.sep_layer = 0 # teacher sep at 0
logger.info('supervise decompose layer: {}'.format(FLAGS.sep_layer))
with tf.variable_scope("teacher", reuse=tf.AUTO_REUSE):
teacher_outputs = get_decomposed_qa_outputs(
FLAGS, features, is_training)
else:
teacher_outputs = None
if FLAGS.decompose:
FLAGS.sep_layer = sep_layer
logger.info('decompose at layer: {}'.format(FLAGS.sep_layer))
outputs = get_decomposed_qa_outputs(FLAGS, features, is_training)
else:
logger.info('running in normal mode')
outputs = get_qa_outputs(FLAGS, features, is_training)
# ### Check model parameters
num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
logger.info('#params: {}'.format(num_params))
scaffold_fn = None
# ### Evaluation mode
if mode == tf.estimator.ModeKeys.PREDICT:
if FLAGS.init_checkpoint:
logger.info(
"init_checkpoint not being used in predict mode.")
predictions = {
"feature_id": features["feature_id"],
"start_logits": outputs["start_logits"],
"end_logits": outputs["end_logits"],
"cls_logits": outputs["cls_logits"]
}
if FLAGS.use_tpu:
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)
else:
output_spec = tf.estimator.EstimatorSpec(
mode=mode, predictions=predictions)
return output_spec
# ## Compute loss
seq_length = FLAGS.max_seq_length
def compute_loss(log_probs, positions, depth=seq_length):
one_hot_positions = tf.one_hot(
positions, depth=depth, dtype=tf.float32)
loss = - tf.reduce_sum(one_hot_positions * log_probs, axis=-1)
loss = tf.reduce_mean(loss)
return loss
start_loss = compute_loss(
outputs["start_log_probs"], features["answer_start"])
end_loss = compute_loss(
outputs["end_log_probs"], features["answer_end"])
total_loss = (start_loss + end_loss) * 0.5
cls_loss = compute_loss(
outputs["cls_log_probs"], features["cls"], depth=FLAGS.num_classes)
# note(zhiliny): by default multiply the loss by 0.5 so that
# the scale is comparable to start_loss and end_loss
total_loss += cls_loss * 0.5
monitor_dict = {"loss/start": start_loss,
"loss/end": end_loss,
"loss/cls": cls_loss,
'loss/ce': total_loss}
if teacher_outputs is not None:
ce_loss = total_loss
gamma = FLAGS.ll_gamma
alpha = FLAGS.dl_alpha
beta = FLAGS.ul_beta
# supervise upper and logits
temp = FLAGS.temperature
kd_loss = get_distill_loss(teacher_outputs, outputs, temp)
mse_loss = get_upper_loss(teacher_outputs, outputs)
monitor_dict["loss/kd"] = kd_loss
monitor_dict["loss/mse"] = mse_loss
total_loss = gamma * ce_loss + alpha * kd_loss + beta * mse_loss
# ### Configuring the optimizer
all_trainable_variables = tf.trainable_variables()
# set fine tune scope
if FLAGS.tune_scopes:
tune_scopes = FLAGS.tune_scopes.split(',')
else:
tune_scopes = None
logger.info('tune_scopes: {}'.format(tune_scopes))
if isinstance(tune_scopes, list):
scoped_variables = []
for scope in tune_scopes:
scoped_variables.extend(tf.trainable_variables(scope))
trainable_variables = scoped_variables
else:
trainable_variables = all_trainable_variables
if FLAGS.init_scopes:
init_scopes = FLAGS.init_scopes.split(',')
else:
init_scopes = None
logger.info('init_scopes: {}'.format(init_scopes))
if isinstance(init_scopes, list):
to_be_init_variables = []
for scope in init_scopes:
to_be_init_variables.extend(tf.trainable_variables(scope))
else:
to_be_init_variables = all_trainable_variables
initialized_variable_names = {}
scaffold_fn = None
# ### load pretrained models
init_checkpoint = FLAGS.init_checkpoint
if init_checkpoint:
logger.info("Initialize from the ckpt {}".format(init_checkpoint))
assign_map, initialized_variable_names = my_init_from_checkpoint(
init_checkpoint, to_be_init_variables)
# logger.info('assign_map: \n{}'.format(assign_map))
# logger.info('initialized_variable_names: \n{}'.format(
# initialized_variable_names))
if FLAGS.use_tpu:
def tpu_scaffold():
tf.train.init_from_checkpoint(init_checkpoint, assign_map)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
tf.train.init_from_checkpoint(init_checkpoint, assign_map)
logger.info("**** Initialized Variables ****")
for var in to_be_init_variables:
init_str = ""
if var.name in initialized_variable_names:
init_str = ", *INIT*"
logger.info(" name=%s, shape=%s%s",
var.name, var.shape, init_str)
if mode == tf.estimator.ModeKeys.TRAIN:
logger.info("**** Trainable Variables ****")
for var in trainable_variables:
init_str = ""
if var.name in initialized_variable_names:
init_str = ", *INIT_AND_TRAINABLE*"
logger.info("*TRAINABLE* name=%s, shape=%s%s",
var.name, var.shape, init_str)
train_op, learning_rate, _ = get_train_op(
FLAGS, total_loss, trainable_variables=trainable_variables)
monitor_dict["lr"] = learning_rate
# ### Constucting training TPUEstimatorSpec with new cache.
if FLAGS.use_tpu:
host_call = construct_scalar_host_call(
monitor_dict=monitor_dict,
model_dir=FLAGS.model_dir,
prefix="train/",
reduce_fn=tf.reduce_mean)
train_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode, loss=total_loss, train_op=train_op,
host_call=host_call,
scaffold_fn=scaffold_fn)
else:
train_spec = tf.estimator.EstimatorSpec(
mode=mode, loss=total_loss, train_op=train_op)
return train_spec
def summarize_sequence(summary_type,
hidden,
d_model,
n_head,
d_head,
dropout,
dropatt,
input_mask,
is_training,
initializer,
scope=None,
reuse=None,
use_proj=True):
"""
Different classification tasks may not may not share the same parameters
to summarize the sequence features.
If shared, one can keep the `scope` to the default value `None`.
Otherwise, one should specify a different `scope` for each task.
"""
with tf.variable_scope(scope, 'sequnece_summary', reuse=reuse):
if summary_type == 'last':
summary = hidden[-1]
elif summary_type == 'first':
summary = hidden[0]
elif summary_type == 'mean':
summary = tf.reduce_mean(hidden, axis=0)
elif summary_type == 'attn':
bsz = tf.shape(hidden)[1]
summary_bias = tf.get_variable('summary_bias', [d_model],
dtype=hidden.dtype,
initializer=initializer)
summary_bias = tf.tile(summary_bias[None, None], [1, bsz, 1])
if input_mask is not None:
input_mask = input_mask[None, :, :, None]
summary = multihead_attn(summary_bias,
hidden,
hidden,
input_mask,
d_model,
n_head,
d_head,
dropout,
dropatt,
is_training,
initializer,
residual=False)
summary = summary[0]
else:
raise ValueError(
'Unsupported summary type {}'.format(summary_type))
# use another projection as in BERT
if use_proj:
summary = tf.layers.dense(summary,
d_model,
activation=tf.tanh,
kernel_initializer=initializer,
name='summary')
# dropout
summary = tf.layers.dropout(summary,
dropout,
training=is_training,
name='dropout')
return summary
def transformer_xl(
input_ids,
n_token,
n_layer,
d_model,
n_head,
d_head,
d_inner,
dropout,
dropatt,
attn_type,
is_training,
initializer,
# bi_data, mem_len=None,
# inp_q=None, mems=None, same_length=False,
clamp_len=-1,
untie_r=False,
use_tpu=True,
input_mask=None,
seg_id=None,
# perm_mask=None, reuse_len=None, target_mapping=None,
ff_activation='relu',
use_bfloat16=False,
scope='transformer',
**kwargs):
"""
Defines a Transformer-XL computation graph with additional
support for XLNet.
Args:
input_ids: int32 Tensor in shape [len, bsz], the input token IDs.
seg_id: int32 Tensor in shape [len, bsz], the input segment IDs.
input_mask: float32 Tensor in shape [len, bsz], the input mask.
0 for real tokens and 1 for padding.
mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory
from previous batches. The length of the list equals n_layer.
If None, no memory is used.
perm_mask: float32 Tensor in shape [len, len, bsz].
If perm_mask[i, j, k] = 0, i attend to j in batch k;
if perm_mask[i, j, k] = 1, i does not attend to j in batch k.
If None, each position attends to all the others.
target_mapping: float32 Tensor in shape [num_predict, len, bsz].
If target_mapping[i, j, k] = 1, the i-th predict in batch k is
on the j-th token.
Only used during pretraining for partial prediction.
Set to None during finetuning.
inp_q: float32 Tensor in shape [len, bsz].
1 for tokens with losses and 0 for tokens without losses.
Only used during pretraining for two-stream attention.
Set to None during finetuning.
n_layer: int, the number of layers.
d_model: int, the hidden size.
n_head: int, the number of attention heads.
d_head: int, the dimension size of each attention head.
d_inner: int, the hidden size in feed-forward layers.
ff_activation: str, "relu" or "gelu".
untie_r: bool, whether to untie the biases in attention.
n_token: int, the vocab size.
is_training: bool, whether in training mode.
use_tpu: bool, whether TPUs are used.
use_bfloat16: bool, use bfloat16 instead of float32.
dropout: float, dropout rate.
dropatt: float, dropout rate on attention probabilities.
init: str, the initialization scheme, either "normal" or "uniform".
init_range: float, initialize the parameters with a uniform distribution
in [-init_range, init_range]. Only effective when init="uniform".
init_std: float, initialize the parameters with a normal distribution
with mean 0 and stddev init_std. Only effective when init="normal".
mem_len: int, the number of tokens to cache.
reuse_len: int, the number of tokens in the currect batch to be cached
and reused in the future.
bi_data: bool, whether to use bidirectional input pipeline.
Usually set to True during pretraining and False during finetuning.
clamp_len: int, clamp all relative distances larger than clamp_len.
-1 means no clamping.
same_length: bool, whether to use the same attention length for each token.
summary_type: str, "last", "first", "mean", or "attn". The method
to pool the input to get a vector representation.
initializer: A tf initializer.
scope: scope name for the computation graph.
"""
# logger.info('memory input {}'.format(mems))
tf_float = tf.bfloat16 if use_bfloat16 else tf.float32
logger.info('Use float type {}'.format(tf_float))
new_mems = []
with tf.variable_scope(scope):
if untie_r:
r_w_bias = tf.get_variable('r_w_bias', [n_layer, n_head, d_head],
dtype=tf_float,
initializer=initializer)
r_r_bias = tf.get_variable('r_r_bias', [n_layer, n_head, d_head],
dtype=tf_float,
initializer=initializer)
else:
r_w_bias = tf.get_variable('r_w_bias', [n_head, d_head],
dtype=tf_float,
initializer=initializer)
r_r_bias = tf.get_variable('r_r_bias', [n_head, d_head],
dtype=tf_float,
initializer=initializer)
batch_size = tf.shape(input_ids)[1]
seq_len = tf.shape(input_ids)[0]
# mlen = tf.shape(mems[0])[0] if mems is not None else 0
mlen = 0
klen = mlen + seq_len
# #### Attention mask
attn_mask = None
# causal attention mask
# if attn_type == 'uni':
# attn_mask = _create_mask(seq_len, mlen, tf_float, same_length)
# attn_mask = attn_mask[:, :, None, None]
# elif attn_type == 'bi':
# attn_mask = None
# else:
# raise ValueError('Unsupported attention type: {}'.format(attn_type))
# data mask: input mask & perm mask
data_mask = input_mask[None]
# if input_mask is not None and perm_mask is not None:
# data_mask = input_mask[None] + perm_mask
# elif input_mask is not None and perm_mask is None:
# data_mask = input_mask[None]
# elif input_mask is None and perm_mask is not None:
# data_mask = perm_mask
# else:
# data_mask = None
if data_mask is not None:
# all mems can be attended to
mems_mask = tf.zeros([tf.shape(data_mask)[0], mlen, batch_size],
dtype=tf_float)
data_mask = tf.concat([mems_mask, data_mask], 1)
if attn_mask is None:
attn_mask = data_mask[:, :, :, None]
else:
attn_mask += data_mask[:, :, :, None]
if attn_mask is not None:
attn_mask = tf.cast(attn_mask > 0, dtype=tf_float)
if attn_mask is not None:
non_tgt_mask = -tf.eye(seq_len, dtype=tf_float)
non_tgt_mask = tf.concat(
[tf.zeros([seq_len, mlen], dtype=tf_float), non_tgt_mask],
axis=-1)
non_tgt_mask = tf.cast(
(attn_mask + non_tgt_mask[:, :, None, None]) > 0,
dtype=tf_float)
else:
non_tgt_mask = None
# #### Word embedding
word_emb_k, lookup_table = embedding_lookup(x=input_ids,
n_token=n_token,
d_embed=d_model,
initializer=initializer,
use_tpu=use_tpu,
dtype=tf_float,
scope='word_embedding')
# if inp_q is not None:
# with tf.variable_scope('mask_emb'):
# mask_emb = tf.get_variable('mask_emb', [1, 1, d_model],
# dtype=tf_float)
# if target_mapping is not None:
# word_emb_q = tf.tile(mask_emb, [tf.shape(target_mapping)[0],
# batch_size, 1])
# else:
# inp_q_ext = inp_q[:, :, None]
# word_emb_q = inp_q_ext * mask_emb + (
# 1 - inp_q_ext) * word_emb_k
output_h = tf.layers.dropout(word_emb_k, dropout, training=is_training)
# if inp_q is not None:
# output_g = tf.layers.dropout(word_emb_q, dropout,
# training=is_training)
# #### Segment embedding
if seg_id is not None:
if untie_r:
r_s_bias = tf.get_variable('r_s_bias',
[n_layer, n_head, d_head],
dtype=tf_float,
initializer=initializer)
else:
# default case (tie)
r_s_bias = tf.get_variable('r_s_bias', [n_head, d_head],
dtype=tf_float,
initializer=initializer)
seg_embed = tf.get_variable('seg_embed',
[n_layer, 2, n_head, d_head],
dtype=tf_float,
initializer=initializer)
# Convert `seg_id` to one-hot `seg_mat`
mem_pad = tf.zeros([mlen, batch_size], dtype=tf.int32)
cat_ids = tf.concat([mem_pad, seg_id], 0)
# `1` indicates not in the same segment [qlen x klen x bsz]
seg_mat = tf.cast(
tf.logical_not(tf.equal(seg_id[:, None], cat_ids[None, :])),
tf.int32)
seg_mat = tf.one_hot(seg_mat, 2, dtype=tf_float)
else:
seg_mat = None
# #### Positional encoding
pos_emb = relative_positional_encoding(seq_len,
klen,
d_model,
clamp_len,
attn_type,
bsz=batch_size,
dtype=tf_float)
pos_emb = tf.layers.dropout(pos_emb, dropout, training=is_training)
# #### Attention layers
# if mems is None:
# mems = [None] * n_layer
mems = [None] * n_layer
for i in range(n_layer):
# cache new mems
# new_mems.append(_cache_mem(output_h, mems[i], mem_len, reuse_len))
new_mems.append(None)
# segment bias
if seg_id is None:
r_s_bias_i = None
seg_embed_i = None
else:
r_s_bias_i = r_s_bias if not untie_r else r_s_bias[i]
seg_embed_i = seg_embed[i]
with tf.variable_scope('layer_{}'.format(i)):
# if inp_q is not None:
# output_h, output_g = two_stream_rel_attn(
# h=output_h,
# g=output_g,
# 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],
# seg_mat=seg_mat,
# r_s_bias=r_s_bias_i,
# seg_embed=seg_embed_i,
# attn_mask_h=non_tgt_mask,
# attn_mask_g=attn_mask,
# mems=mems[i],
# target_mapping=target_mapping,
# d_model=d_model,
# n_head=n_head,
# d_head=d_head,
# dropout=dropout,
# dropatt=dropatt,
# is_training=is_training,
# kernel_initializer=initializer)
# reuse = True
# else:
reuse = False
output_h = rel_multihead_attn(
h=output_h,
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],
seg_mat=seg_mat,
r_s_bias=r_s_bias_i,
seg_embed=seg_embed_i,
attn_mask=non_tgt_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,
reuse=reuse)
# if inp_q is not None:
# output_g = positionwise_ffn(
# inp=output_g,
# d_model=d_model,
# d_inner=d_inner,
# dropout=dropout,
# kernel_initializer=initializer,
# activation_type=ff_activation,
# is_training=is_training)
output_h = positionwise_ffn(inp=output_h,
d_model=d_model,
d_inner=d_inner,
dropout=dropout,
kernel_initializer=initializer,
activation_type=ff_activation,
is_training=is_training,
reuse=reuse)
# if inp_q is not None:
# output = tf.layers.dropout(output_g, dropout, training=is_training)
# else:
# output = tf.layers.dropout(output_h, dropout, training=is_training)
output = tf.layers.dropout(output_h, dropout, training=is_training)
return output, new_mems, lookup_table
def transformer_xl_decomposed(n_token,
n_layer,
d_model,
n_head,
d_head,
d_inner,
dropout,
dropatt,
attn_type,
is_training,
initializer,
q_ids,
ctx_ids,
clamp_len=-1,
untie_r=False,
use_tpu=True,
ff_activation='relu',
use_bfloat16=False,
sep_layer=9,
q_attn_mask=None,
c_attn_mask=None,
qc_attn_mask=None,
q_seq_len=None,
ctx_seq_len=None,
scope='transformer',
**kwargs):
tf_float = tf.bfloat16 if use_bfloat16 else tf.float32
logger.info('Use float type {}'.format(tf_float))
# new_mems = []
with tf.variable_scope(scope):
if untie_r:
r_w_bias = tf.get_variable('r_w_bias', [n_layer, n_head, d_head],
dtype=tf_float,
initializer=initializer)
r_r_bias = tf.get_variable('r_r_bias', [n_layer, n_head, d_head],
dtype=tf_float,
initializer=initializer)
else:
r_w_bias = tf.get_variable('r_w_bias', [n_head, d_head],
dtype=tf_float,
initializer=initializer)
r_r_bias = tf.get_variable('r_r_bias', [n_head, d_head],
dtype=tf_float,
initializer=initializer)
# batch_size = tf.shape(input_ids)[1]
# seq_len = tf.shape(input_ids)[0]
batch_size = tf.shape(q_ids)[1]
# mlen = tf.shape(mems[0])[0] if mems is not None else 0
# mlen = 0
# klen = mlen + seq_len
# #### Attention mask
attn_mask = None
# data_mask = input_mask[None]
# if data_mask is not None:
# all mems can be attended to
# mems_mask = tf.zeros([tf.shape(data_mask)[0], mlen, batch_size],
# dtype=tf_float)
# data_mask = tf.concat([mems_mask, data_mask], 1)
# if attn_mask is None:
# attn_mask = data_mask[:, :, :, None]
# else:
# attn_mask += data_mask[:, :, :, None]
# non_tgt_mask = None
# #### Word embedding
q_emb, lookup_table = embedding_lookup(x=q_ids,
n_token=n_token,
d_embed=d_model,
initializer=initializer,
use_tpu=use_tpu,
dtype=tf_float,
scope='word_embedding')
c_emb, _ = embedding_lookup(x=ctx_ids,
n_token=n_token,
d_embed=d_model,
initializer=initializer,
use_tpu=use_tpu,
dtype=tf_float,
reuse=True,
scope='word_embedding')
q_output_h = tf.layers.dropout(q_emb, dropout, training=is_training)
ctx_output_h = tf.layers.dropout(c_emb, dropout, training=is_training)
# #### Segment embedding
if untie_r:
r_s_bias = tf.get_variable('r_s_bias', [n_layer, n_head, d_head],
dtype=tf_float,
initializer=initializer)
else:
# default case (tie)
r_s_bias = tf.get_variable('r_s_bias', [n_head, d_head],
dtype=tf_float,
initializer=initializer)
seg_embed = tf.get_variable('seg_embed', [n_layer, 2, n_head, d_head],
dtype=tf_float,
initializer=initializer)
# Convert `seg_id` to one-hot `seg_mat`
# mem_pad = tf.zeros([mlen, batch_size], dtype=tf.int32)
# cat_ids = tf.concat([mem_pad, seg_id], 0)
# `1` indicates not in the same segment [qlen x klen x bsz]
ctx_seg_ids = tf.zeros_like(ctx_ids, dtype=tf.int32)
ctx_seg_mat = tf.cast(
tf.logical_not(tf.equal(ctx_seg_ids[:, None],
ctx_seg_ids[None, :])), tf.int32)
ctx_seg_mat = tf.one_hot(ctx_seg_mat, 2, dtype=tf_float)
q_seg_ids = tf.ones_like(q_ids, dtype=tf.int32)
q_seg_mat = tf.cast(
tf.logical_not(tf.equal(q_seg_ids[:, None], q_seg_ids[None, :])),
tf.int32)
q_seg_mat = tf.one_hot(q_seg_mat, 2, dtype=tf_float)
seg_ids = tf.concat([ctx_seg_ids, q_seg_ids], axis=0)
seg_mat = tf.cast(
tf.logical_not(tf.equal(seg_ids[:, None], seg_ids[None, :])),
tf.int32)
seg_mat = tf.one_hot(seg_mat, 2, dtype=tf_float)
# #### Positional encoding FIXME: better use of relative pos emb
q_pos_emb = relative_positional_encoding(q_seq_len,
q_seq_len,
d_model,
clamp_len,
attn_type,
bsz=batch_size,
dtype=tf_float)
q_pos_emb = tf.layers.dropout(q_pos_emb, dropout, training=is_training)
ctx_pos_emb = relative_positional_encoding(ctx_seq_len,
ctx_seq_len,
d_model,
clamp_len,
attn_type,
bsz=batch_size,
dtype=tf_float)
ctx_pos_emb = tf.layers.dropout(ctx_pos_emb,
dropout,
training=is_training)
# pos_emb = tf.concat([ctx_pos_emb, q_pos_emb], axis=0)
seq_len = ctx_seq_len + q_seq_len
pos_emb = relative_positional_encoding(seq_len,
seq_len,
d_model,
clamp_len,
attn_type,
bsz=batch_size,
dtype=tf_float)
pos_emb = tf.layers.dropout(pos_emb, dropout, training=is_training)
# ctx_pos_emb = pos_emb[q_seq_len:q_seq_len + 2 * ctx_seq_len, :, :]
# q_pos_emb1 = pos_emb[:q_seq_len, :, :]
# q_pos_emb2 = pos_emb[q_seq_len + 2 * ctx_seq_len:, :, :]
# q_pos_emb = tf.concat([q_pos_emb1, q_pos_emb2], axis=0)
# #### Attention layers
# mems = [None] * n_layer
for i in range(sep_layer):
r_s_bias_i = r_s_bias if not untie_r else r_s_bias[i]
r_w_bias_i = r_w_bias if not untie_r else r_w_bias[i]
r_r_bias_i = r_r_bias if not untie_r else r_r_bias[i]
seg_embed_i = seg_embed[i]
with tf.variable_scope('layer_{}'.format(i)):
ctx_output_h = rel_multihead_attn(
h=ctx_output_h,
r=ctx_pos_emb,
r_w_bias=r_w_bias_i,
r_r_bias=r_r_bias_i,
r_s_bias=r_s_bias_i,
seg_mat=ctx_seg_mat,
seg_embed=seg_embed_i,
attn_mask=c_attn_mask,
mems=None,
d_model=d_model,
n_head=n_head,
d_head=d_head,
dropout=dropout,
dropatt=dropatt,
is_training=is_training,
kernel_initializer=initializer,
reuse=False)
ctx_output_h = positionwise_ffn(inp=ctx_output_h,
d_model=d_model,
d_inner=d_inner,
dropout=dropout,
kernel_initializer=initializer,
activation_type=ff_activation,
is_training=is_training,
reuse=False)
q_output_h = rel_multihead_attn(h=q_output_h,
r=q_pos_emb,
r_w_bias=r_w_bias_i,
r_r_bias=r_r_bias_i,
r_s_bias=r_s_bias_i,
seg_mat=q_seg_mat,
seg_embed=seg_embed_i,
attn_mask=q_attn_mask,
mems=None,
d_model=d_model,
n_head=n_head,
d_head=d_head,
dropout=dropout,
dropatt=dropatt,
is_training=is_training,
kernel_initializer=initializer,
reuse=tf.AUTO_REUSE)
q_output_h = positionwise_ffn(inp=q_output_h,
d_model=d_model,
d_inner=d_inner,
dropout=dropout,
kernel_initializer=initializer,
activation_type=ff_activation,
is_training=is_training,
reuse=tf.AUTO_REUSE)
# concat all q, ctx related variables
output_h = tf.concat([ctx_output_h, q_output_h], axis=0)
upper_outputs = []
for i in range(sep_layer, n_layer):
r_s_bias_i = r_s_bias if not untie_r else r_s_bias[i]
r_w_bias_i = r_w_bias if not untie_r else r_w_bias[i]
r_r_bias_i = r_r_bias if not untie_r else r_r_bias[i]
seg_embed_i = seg_embed[i]
with tf.variable_scope('layer_{}'.format(i)):
output_h = rel_multihead_attn(h=output_h,
r=pos_emb,
seg_mat=seg_mat,
r_w_bias=r_w_bias_i,
r_r_bias=r_r_bias_i,
r_s_bias=r_s_bias_i,
seg_embed=seg_embed_i,
attn_mask=qc_attn_mask,
mems=None,
d_model=d_model,
n_head=n_head,
d_head=d_head,
dropout=dropout,
dropatt=dropatt,
is_training=is_training,
kernel_initializer=initializer,
reuse=False)
output_h = positionwise_ffn(inp=output_h,
d_model=d_model,
d_inner=d_inner,
dropout=dropout,
kernel_initializer=initializer,
activation_type=ff_activation,
is_training=is_training,
reuse=False)
upper_outputs.append(output_h)
output = tf.layers.dropout(output_h, dropout, training=is_training)
upper_outputs[-1] = output
return upper_outputs
def get_decomposed_qa_outputs(FLAGS, features, is_training):
question_ids = features["question_ids"]
context_ids = features["context_ids"]
seq_len = FLAGS.max_seq_length
q_seq_len = FLAGS.max_first_length + 2
ctx_seq_len = seq_len - q_seq_len
q_mask_int = tf.cast(tf.cast(question_ids, tf.bool), tf.int32)
cls_index = tf.reshape(
tf.reduce_sum(q_mask_int, axis=1) + ctx_seq_len, [-1])
# 0 for mask out
# q_zeros = tf.zeros_like(question_ids)
# p_ids = tf.concat([context_ids, q_zeros], axis=1)
# p_mask = tf.cast(tf.cast(p_ids, tf.bool), tf.float32)
question_ids = tf.transpose(question_ids, [1, 0])
context_ids = tf.transpose(context_ids, [1, 0])
q_attn_mask = get_attention_mask(question_ids, q_seq_len)
c_attn_mask = get_attention_mask(context_ids, ctx_seq_len)
qc_attn_mask = get_attention_mask(
tf.concat([context_ids, question_ids], axis=0), seq_len)
xlnet_config = xlnet.XLNetConfig(json_path=FLAGS.model_config_path)
run_config = xlnet.create_run_config(is_training, True, FLAGS)
initializer = xlnet._get_initializer(run_config)
tfm_args = dict(
n_token=xlnet_config.n_token,
initializer=initializer,
attn_type="bi",
n_layer=xlnet_config.n_layer,
d_model=xlnet_config.d_model,
n_head=xlnet_config.n_head,
d_head=xlnet_config.d_head,
d_inner=xlnet_config.d_inner,
ff_activation=xlnet_config.ff_activation,
untie_r=xlnet_config.untie_r,
is_training=run_config.is_training,
use_bfloat16=run_config.use_bfloat16,
use_tpu=run_config.use_tpu,
dropout=run_config.dropout,
dropatt=run_config.dropatt,
# mem_len=run_config.mem_len,
# reuse_len=run_config.reuse_len,
# bi_data=run_config.bi_data,
clamp_len=run_config.clamp_len,
# same_length=run_config.same_length,
ctx_ids=context_ids,
q_ids=question_ids,
q_seq_len=q_seq_len,
ctx_seq_len=ctx_seq_len,
sep_layer=FLAGS.sep_layer,
q_attn_mask=q_attn_mask,
c_attn_mask=c_attn_mask,
qc_attn_mask=qc_attn_mask,
)
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
upper_outputs = transformer_xl_decomposed(**tfm_args)
output = upper_outputs[-1]
return_dict = {'upper_outputs': upper_outputs}
with tf.variable_scope("logits"):
# logits: seq, batch_size, 2
logits = tf.layers.dense(output, 2, kernel_initializer=initializer)
# logits: 2, batch_size, seq
logits = tf.transpose(logits, [2, 1, 0])
# start_logits: batch_size, seq
# end_logits: batch_size, seq
start_logits, end_logits = tf.unstack(logits, axis=0)
# start_logits_masked = start_logits * p_mask - 1e30 * (1 - p_mask)
# start_log_probs = tf.nn.log_softmax(start_logits_masked, -1)
start_log_probs = tf.nn.log_softmax(start_logits, -1)
# end_logits_masked = end_logits * p_mask - 1e30 * (1 - p_mask)
# end_log_probs = tf.nn.log_softmax(end_logits_masked, -1)
end_log_probs = tf.nn.log_softmax(end_logits, -1)
return_dict["start_logits"] = start_logits
return_dict["end_logits"] = end_logits
if is_training:
return_dict["start_log_probs"] = start_log_probs
return_dict["end_log_probs"] = end_log_probs
# an additional layer to predict answer class, 0: span, 1:yes, 2:no
with tf.variable_scope("answer_class"):
# get the representation of CLS
cls_index = tf.one_hot(cls_index, seq_len, axis=-1, dtype=tf.float32)
cls_feature = tf.einsum("lbh,bl->bh", output, cls_index)
ans_feature = tf.layers.dense(cls_feature,
xlnet_config.d_model,
activation=tf.tanh,
kernel_initializer=initializer,
name='pooler')
ans_feature = tf.layers.dropout(ans_feature,
FLAGS.dropout,
training=is_training)
# hotpot has 3 classes,
# squad 2.0 has 2 classes
cls_logits = tf.layers.dense(ans_feature,
FLAGS.num_classes,
kernel_initializer=initializer,
name="cls")
cls_log_probs = tf.nn.log_softmax(cls_logits, -1)
return_dict["cls_logits"] = cls_logits
if is_training:
return_dict["cls_log_probs"] = cls_log_probs
return return_dict
def get_qa_outputs(FLAGS, features, is_training):
"""Loss for downstream span-extraction QA tasks such as SQuAD."""
input_ids = features["input_ids"]
seg_id = features["segment_ids"]
input_mask_int = tf.cast(tf.cast(input_ids, tf.bool), tf.int32)
cls_index = tf.reshape(tf.reduce_sum(input_mask_int, axis=1), [-1])
p_mask = tf.cast(tf.cast(seg_id, tf.bool), tf.float32)
input_ids = tf.transpose(input_ids, [1, 0])
input_mask = 1 - tf.cast(input_mask_int, tf.float32)
input_mask = tf.transpose(input_mask, [1, 0])
seg_id = tf.transpose(seg_id, [1, 0])
seq_len = tf.shape(input_ids)[0]
xlnet_config = xlnet.XLNetConfig(json_path=FLAGS.model_config_path)
run_config = xlnet.create_run_config(is_training, True, FLAGS)
xlnet_model = xlnet.XLNetModel(
xlnet_config=xlnet_config,
run_config=run_config,
input_ids=input_ids,
seg_ids=seg_id,
input_mask=input_mask)
output = xlnet_model.get_sequence_output()
initializer = xlnet_model.get_initializer()
return_dict = {}
with tf.variable_scope("logits"):
# logits: seq, batch_size, 2
logits = tf.layers.dense(output, 2, kernel_initializer=initializer)
# logits: 2, batch_size, seq
logits = tf.transpose(logits, [2, 1, 0])
# start_logits: batch_size, seq
# end_logits: batch_size, seq
start_logits, end_logits = tf.unstack(logits, axis=0)
start_logits_masked = start_logits * (1 - p_mask) - 1e30 * p_mask
start_log_probs = tf.nn.log_softmax(start_logits_masked, -1)
end_logits_masked = end_logits * (1 - p_mask) - 1e30 * p_mask
end_log_probs = tf.nn.log_softmax(end_logits_masked, -1)
if is_training:
return_dict["start_log_probs"] = start_log_probs
return_dict["end_log_probs"] = end_log_probs
else:
return_dict["start_logits"] = start_logits
return_dict["end_logits"] = end_logits
# an additional layer to predict answer class, 0: span, 1:yes, 2:no
with tf.variable_scope("answer_class"):
# get the representation of CLS
cls_index = tf.one_hot(cls_index, seq_len, axis=-1, dtype=tf.float32)
cls_feature = tf.einsum("lbh,bl->bh", output, cls_index)
ans_feature = tf.layers.dense(cls_feature, xlnet_config.d_model,
activation=tf.tanh,
kernel_initializer=initializer,
name='pooler')
ans_feature = tf.layers.dropout(ans_feature, FLAGS.dropout,
training=is_training)
# hotpot has 3 classes,
# squad 2.0 has 2 classes
cls_logits = tf.layers.dense(ans_feature, FLAGS.num_classes,
kernel_initializer=initializer,
name="cls")
cls_log_probs = tf.nn.log_softmax(cls_logits, -1)
if is_training:
return_dict["cls_log_probs"] = cls_log_probs
return_dict["cls_logits"] = cls_logits
return return_dict
def __init__(self, xlnet_config, run_config, input_ids, seg_ids, input_mask,
# mems=None, perm_mask=None, target_mapping=None, inp_q=None,
**kwargs):
"""
Args:
xlnet_config: XLNetConfig,
run_config: RunConfig,
input_ids: int32 Tensor in shape [len, bsz], the input token IDs.
seg_ids: int32 Tensor in shape [len, bsz], the input segment IDs.
input_mask: float32 Tensor in shape [len, bsz], the input mask.
0 for real tokens and 1 for padding.
mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory
from previous batches. The length of the list equals n_layer.
If None, no memory is used.
perm_mask: float32 Tensor in shape [len, len, bsz].
If perm_mask[i, j, k] = 0, i attend to j in batch k;
if perm_mask[i, j, k] = 1, i does not attend to j in batch k.
If None, each position attends to all the others.
target_mapping: float32 Tensor in shape [num_predict, len, bsz].
If target_mapping[i, j, k] = 1, the i-th predict in batch k is
on the j-th token.
Only used during pretraining for partial prediction.
Set to None during finetuning.
inp_q: float32 Tensor in shape [len, bsz].
1 for tokens with losses and 0 for tokens without losses.
Only used during pretraining for two-stream attention.
Set to None during finetuning.
"""
initializer = _get_initializer(run_config)
tfm_args = dict(
n_token=xlnet_config.n_token,
initializer=initializer,
attn_type="bi",
n_layer=xlnet_config.n_layer,
d_model=xlnet_config.d_model,
n_head=xlnet_config.n_head,
d_head=xlnet_config.d_head,
d_inner=xlnet_config.d_inner,
ff_activation=xlnet_config.ff_activation,
untie_r=xlnet_config.untie_r,
is_training=run_config.is_training,
use_bfloat16=run_config.use_bfloat16,
use_tpu=run_config.use_tpu,
dropout=run_config.dropout,
dropatt=run_config.dropatt,
# mem_len=run_config.mem_len,
# reuse_len=run_config.reuse_len,
# bi_data=run_config.bi_data,
clamp_len=run_config.clamp_len,
# same_length=run_config.same_length
)
input_args = dict(
input_ids=input_ids,
seg_id=seg_ids,
input_mask=input_mask,
# mems=mems,
# perm_mask=perm_mask,
# target_mapping=target_mapping,
# inp_q=inp_q,
)
tfm_args.update(input_args)
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
(self.output, self.new_mems, self.lookup_table
) = modeling.transformer_xl(**tfm_args)
self.input_mask = input_mask
self.initializer = initializer
self.xlnet_config = xlnet_config
self.run_config = run_config
