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