def mlm_loss(features, labels, mems, n_token, is_training): """Standard MLM loss as in BERT.""" del labels del mems initializer = _get_initializer() #### Unpack input masked_inp = features["masked_input"] type_id = features["type_id"] target_mapping = features["target_mapping"] target_mask = features["target_mask"] target = features["target"] monitor_dict = {} #### Transformer Model with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) masked_embed, word_embed_table = inp_func(inputs=masked_inp, type_id=type_id, return_embed_table=True) tfm_func = _get_tfm_func(initializer, is_training, phase="pretrain") output, _ = tfm_func(inputs=masked_embed, input_mask=None, perm_mask=None) lm_loss, _ = model.lm_loss(hidden=output, target=target, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=word_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=target_mapping, return_logits=False, use_tpu=FLAGS.use_tpu) if lm_loss.dtype != tf.float32: lm_loss = tf.cast(lm_loss, tf.float32) if FLAGS.sample_strategy == "token_span": if target_mask.dtype != tf.float32: target_mask = tf.cast(target_mask, tf.float32) total_loss = (tf.reduce_sum(lm_loss * target_mask) / tf.reduce_sum(target_mask)) else: total_loss = tf.reduce_mean(lm_loss) monitor_dict["lm_loss"] = total_loss return total_loss, {}, monitor_dict
def get_lm_loss(features, mems, n_token, is_training): """LM loss.""" del mems initializer = _get_initializer() #### Unpack input inputs = features["inputs"] target = features["target"] type_id = features["type_id"] monitor_dict = {} #### Transformer Model with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) input_embed, word_embed_table = inp_func(inputs=inputs, type_id=type_id, return_embed_table=True) tfm_func = _get_tfm_func(initializer, is_training, phase="pretrain") output, _ = tfm_func(inputs=input_embed, input_mask=None, perm_mask=None, causal=True) lm_loss, _ = model.lm_loss(hidden=output, target=target, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=word_embed_table, tie_weight=FLAGS.tie_weight, return_logits=False, use_tpu=FLAGS.use_tpu) if lm_loss.dtype != tf.float32: lm_loss = tf.cast(lm_loss, tf.float32) total_loss = tf.reduce_mean(lm_loss) monitor_dict["lm_loss"] = total_loss return total_loss, {}, monitor_dict
def dae_loss(features, labels, mems, n_token, is_training): """DAE loss with generator.""" del mems del labels initializer = _get_initializer() monitor_dict = {} ##### unpack input gen_inp = features["gen_inp"] gen_tgt = features["gen_tgt"] gen_mask_map = features["gen_mask_map"] gen_tgt_mask = features["gen_tgt_mask"] enc_mask = features["enc_mask"] enc_type = features["enc_type"] dec_inp = features["dec_inp"] dec_tgt = features["dec_tgt"] dec_mask = features["dec_mask"] dec_type = features["dec_type"] edit_label = features["edit_label"] dec_mask_map = features["dec_mask_map"] dec_masked_tgt = features["dec_masked_tgt"] dec_lm_tgt_mask = features["dec_lm_tgt_mask"] rep_enc2dec_full = features["rep_enc2dec_full"] rep_enc2dec_part = features["rep_enc2dec_part"] #### Shared input embedding (for generator) with tf.variable_scope("input", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) gen_embed, shared_embed_table = inp_func(inputs=gen_inp, type_id=enc_type, return_embed_table=True) #### Generator TFM with tf.variable_scope("generator", reuse=tf.AUTO_REUSE): gen_tfm_func = _get_tfm_func(initializer, is_training, "pretrain", shrink=FLAGS.gen_shrink) gen_output, _ = gen_tfm_func(inputs=gen_embed, input_mask=enc_mask, perm_mask=None) gen_lm_loss, _, logits = model.lm_loss(hidden=gen_output, target=gen_tgt, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=shared_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=gen_mask_map, return_logits=True, use_tpu=FLAGS.use_tpu) if gen_lm_loss.dtype != tf.float32: gen_lm_loss = tf.cast(gen_lm_loss, tf.float32) gen_tgt_mask = tf.cast(gen_tgt_mask, gen_lm_loss.dtype) gen_loss = (tf.reduce_sum(gen_lm_loss * gen_tgt_mask) / tf.reduce_sum(gen_tgt_mask)) monitor_dict["gen_loss"] = gen_loss total_loss = gen_loss #### Sample from generator uniform = tf.random.uniform(minval=0, maxval=1, shape=logits.shape, dtype=tf.float32) gumbel = -tf.log(-tf.log(uniform + 1e-9) + 1e-9) samples = tf.argmax(logits + gumbel, -1) gen_tokens = tf.cast(samples, tf.int32) # map `num_predict` samples to full length samples = tf.einsum("bm,bml->bl", tf.cast(samples, tf.float32), tf.cast(gen_mask_map, tf.float32)) samples = tf.cast(samples, tf.int32) enc_inp = tf.where(tf.equal(gen_inp, FLAGS.mask_id), samples, gen_inp) #### get the mask for generated same token as the target same_mask = tf.equal(gen_tokens, gen_tgt) same_mask = (tf.cast(same_mask, rep_enc2dec_full.dtype) * tf.cast(gen_tgt_mask, rep_enc2dec_full.dtype)) # monitor how many generated tokens are the same as the real ones same_prec = (tf.reduce_sum(tf.cast(same_mask, gen_tgt_mask.dtype)) / tf.reduce_sum(gen_tgt_mask)) monitor_dict["same_percent"] = same_prec # If same, change the edit_label to original (0) dec_same_mask_full = tf.einsum("bi,bil->bl", same_mask, rep_enc2dec_full) edit_label = tf.where( tf.cast(dec_same_mask_full, tf.bool), tf.zeros(tf.shape(dec_same_mask_full), dtype=edit_label.dtype), edit_label) # If same, exclude from LM loss dec_same_mask_part = tf.einsum("bi,bij->bj", same_mask, rep_enc2dec_part) dec_diff_mask_part = 1.0 - tf.cast(dec_same_mask_part, dec_lm_tgt_mask.dtype) dec_lm_tgt_mask = dec_lm_tgt_mask * dec_diff_mask_part #### Shared input embedding (for encoder) with tf.variable_scope("input", reuse=tf.AUTO_REUSE): enc_embed = inp_func(inputs=enc_inp, type_id=enc_type) #### Encoder TFM with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE): enc_tfm_func = _get_tfm_func(initializer, is_training, "pretrain") enc_output, _ = enc_tfm_func(inputs=enc_embed, input_mask=enc_mask, perm_mask=None) #### Shared input embedding (for decoder) with tf.variable_scope("input", reuse=tf.AUTO_REUSE): dec_embed = inp_func(inputs=dec_inp, type_id=dec_type) #### Decoder TFM with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE): dec_tfm_func = _get_tfm_func(initializer, is_training, "pretrain", causal=True) dec_output, _ = dec_tfm_func(inputs=dec_embed, input_mask=dec_mask, context=enc_output, context_mask=enc_mask[:, None, :, None], perm_mask=None) #### edit type loss edit_loss, edit_logits = model.cls_loss(hidden=dec_output, target=edit_label, n_cls=4, d_model=FLAGS.d_model, initializer=initializer, target_mapping=None, hidden_mapping=None, return_logits=True, scope="edit_type_loss") dec_tgt_mask = tf.cast(1.0 - dec_mask, tf.float32) edit_loss = (tf.reduce_sum(edit_loss * dec_tgt_mask) / tf.reduce_sum(dec_tgt_mask)) edit_pred = tf.cast(tf.argmax(edit_logits, axis=-1), dtype=edit_label.dtype) edit_corr = tf.cast(tf.equal(edit_pred, edit_label), dtype=tf.float32) edit_accu = (tf.reduce_sum(edit_corr * dec_tgt_mask) / tf.reduce_sum(dec_tgt_mask)) if FLAGS.edit_weight > 0: # monitor monitor_dict["edit_loss"] = edit_loss monitor_dict["accu_edit"] = edit_accu def get_class_acc(label_id): mask = tf.ones(tf.shape(edit_label), dtype=edit_label.dtype) * label_id mask = tf.equal(edit_label, mask) mask = tf.cast(mask, dtype=dec_tgt_mask.dtype) * dec_tgt_mask accu = tf.reduce_sum(edit_corr * mask) / tf.reduce_sum(mask) return accu monitor_dict["accu_orig"] = get_class_acc(0) if FLAGS.del_ratio > 0: monitor_dict["accu_del"] = get_class_acc(FLAGS.del_label) if FLAGS.ins_ratio > 0: monitor_dict["accu_ins"] = get_class_acc(FLAGS.ins_label) if FLAGS.rep_ratio > 0: monitor_dict["accu_rep"] = get_class_acc(FLAGS.rep_label) # accumulate total loss total_loss += FLAGS.edit_weight * edit_loss #### next-token prediction loss if FLAGS.mask_edited_only: lm_loss, _ = model.lm_loss(hidden=dec_output, target=dec_masked_tgt, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=shared_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=dec_mask_map, use_tpu=FLAGS.use_tpu) if lm_loss.dtype != tf.float32: lm_loss = tf.cast(lm_loss, tf.float32) dec_lm_tgt_mask = tf.cast(dec_lm_tgt_mask, lm_loss.dtype) lm_loss = (tf.reduce_sum(lm_loss * dec_lm_tgt_mask) / tf.reduce_sum(dec_lm_tgt_mask)) else: lm_loss, _ = model.lm_loss(hidden=dec_output, target=dec_tgt, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=shared_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=None, use_tpu=FLAGS.use_tpu) lm_loss = (tf.reduce_sum(lm_loss * dec_tgt_mask) / tf.reduce_sum(dec_tgt_mask)) if FLAGS.lm_weight > 0: # monitor monitor_dict["lm_loss"] = lm_loss # accumulate total loss total_loss += FLAGS.lm_weight * lm_loss return total_loss, {}, monitor_dict
def dae_joint_loss(features, labels, mems, n_token, is_training): """DAE loss with generator.""" del mems del labels initializer = _get_initializer() monitor_dict = {} ##### unpack input gen_inp = features["gen_inp"] gen_tgt = features["gen_tgt"] gen_mask_map = features["gen_mask_map"] gen_tgt_mask = features["gen_tgt_mask"] enc_mask = features["enc_mask"] enc_type = features["enc_type"] # enc_edit_label = features["enc_edit_label"] dec_inp = features["dec_inp"] dec_tgt = features["dec_tgt"] dec_mask = features["dec_mask"] dec_type = features["dec_type"] edit_label = features["edit_label"] dec_mask_map = features["dec_mask_map"] dec_masked_tgt = features["dec_masked_tgt"] dec_lm_tgt_mask = features["dec_lm_tgt_mask"] rep_enc2dec_full = features["rep_enc2dec_full"] rep_enc2dec_part = features["rep_enc2dec_part"] enc_pos = features["enc_pos"] dec_pos = features["dec_pos"] if FLAGS.double_type: # offer a indicator to differeniate encoder and decoder dec_type = dec_type + FLAGS.n_type initializer = _get_initializer() if FLAGS.use_bfloat16: tf_float = tf.bfloat16 else: tf_float = tf.float32 #### Shared input embedding (for generator) with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) gen_embed, shared_embed_table = inp_func(inputs=gen_inp, type_id=enc_type, return_embed_table=True) #### Generator TFM with tf.variable_scope("generator", reuse=tf.AUTO_REUSE): gen_tfm_func = _get_tfm_func(initializer, is_training, "pretrain", shrink=FLAGS.gen_shrink) gen_output, _ = gen_tfm_func(inputs=gen_embed, input_mask=enc_mask, perm_mask=None) gen_lm_loss, _, logits = model.lm_loss(hidden=gen_output, target=gen_tgt, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=shared_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=gen_mask_map, return_logits=True, use_tpu=FLAGS.use_tpu) if gen_lm_loss.dtype != tf.float32: gen_lm_loss = tf.cast(gen_lm_loss, tf.float32) gen_tgt_mask = tf.cast(gen_tgt_mask, gen_lm_loss.dtype) gen_loss = (tf.reduce_sum(gen_lm_loss * gen_tgt_mask) / tf.reduce_sum(gen_tgt_mask)) monitor_dict["gen_loss"] = gen_loss total_loss = gen_loss #### Sample from generator uniform = tf.random.uniform(minval=0, maxval=1, shape=logits.shape, dtype=tf.float32) gumbel = -tf.log(-tf.log(uniform + 1e-9) + 1e-9) samples = tf.argmax(logits + gumbel, -1) gen_tokens = tf.cast(samples, tf.int32) # map `num_predict` samples to full length samples = tf.einsum("bm,bml->bl", tf.cast(samples, tf.float32), tf.cast(gen_mask_map, tf.float32)) samples = tf.cast(samples, tf.int32) enc_inp = tf.where(tf.equal(gen_inp, FLAGS.mask_id), samples, gen_inp) #### get the mask for generated same token as the target same_mask = tf.equal(gen_tokens, gen_tgt) same_mask = (tf.cast(same_mask, rep_enc2dec_full.dtype) * tf.cast(gen_tgt_mask, rep_enc2dec_full.dtype)) # monitor how many generated tokens are the same as the real ones same_prec = (tf.reduce_sum(tf.cast(same_mask, gen_tgt_mask.dtype)) / tf.reduce_sum(gen_tgt_mask)) monitor_dict["same_percent"] = same_prec # If same, change the edit_label to original (0) dec_same_mask_full = tf.einsum("bi,bil->bl", same_mask, rep_enc2dec_full) edit_label = tf.where( tf.cast(dec_same_mask_full, tf.bool), tf.zeros(tf.shape(dec_same_mask_full), dtype=edit_label.dtype), edit_label) # If same, exclude from LM loss dec_same_mask_part = tf.einsum("bi,bij->bj", same_mask, rep_enc2dec_part) dec_diff_mask_part = 1.0 - tf.cast(dec_same_mask_part, dec_lm_tgt_mask.dtype) dec_lm_tgt_mask = dec_lm_tgt_mask * dec_diff_mask_part # shapes bsz = tf.shape(enc_inp)[0] src_len = tf.shape(enc_inp)[1] tgt_len = tf.shape(dec_inp)[1] ##### format joint model inputs inputs = tf.concat([enc_inp, dec_inp], axis=1) type_id = tf.concat([enc_type, dec_type], axis=1) source_seg = tf.zeros([bsz, src_len], dtype=inputs.dtype) target_seg = tf.ones([bsz, tgt_len], dtype=inputs.dtype) ##### attention mask: note that `1` indicates CANNOT attend # src mask src_to_src = tf.not_equal(source_seg[:, :, None], source_seg[:, None, :]) src_to_tgt = tf.ones([bsz, src_len, tgt_len], dtype=src_to_src.dtype) src_mask = tf.concat([src_to_src, src_to_tgt], axis=2) # tgt mask tgt_to_src = tf.not_equal(target_seg[:, :, None], source_seg[:, None, :]) tgt_to_tgt = tf.not_equal(target_seg[:, :, None], target_seg[:, None, :]) causal_mask = tf.cast(causal_attn_mask(qlen=tgt_len), tgt_to_tgt.dtype) # If any one of them is `1` (indicating cannot attend), i.e. `logical_or`, # then the model should NOT attend tgt_to_tgt = tf.logical_or(tgt_to_tgt, causal_mask) tgt_mask = tf.concat([tgt_to_src, tgt_to_tgt], axis=2) # concat perm_mask = tf.concat([src_mask, tgt_mask], axis=1) perm_mask = tf.cast(perm_mask, tf_float) ##### Position sequence pos_seq = tf.concat([enc_pos, dec_pos], axis=1) #### Transformer Model total_loss = 0 with tf.variable_scope("model", reuse=tf.AUTO_REUSE): input_embed = inp_func(inputs=inputs, pos_seq=pos_seq, type_id=type_id) tfm_func = _get_tfm_func(initializer, is_training, phase="pretrain") output, _ = tfm_func(inputs=input_embed, input_mask=None, perm_mask=perm_mask) #### edit loss tgt_out = output[:, src_len:] edit_loss, edit_logits = model.cls_loss(hidden=tgt_out, target=edit_label, n_cls=4, d_model=FLAGS.d_model, initializer=initializer, target_mapping=None, hidden_mapping=None, return_logits=True, scope="edit_type_loss") dec_tgt_mask = tf.cast(1.0 - dec_mask, tf.float32) edit_loss = (tf.reduce_sum(edit_loss * dec_tgt_mask) / tf.reduce_sum(dec_tgt_mask)) edit_pred = tf.cast(tf.argmax(edit_logits, axis=-1), dtype=edit_label.dtype) edit_corr = tf.cast(tf.equal(edit_pred, edit_label), dtype=tf.float32) edit_accu = (tf.reduce_sum(edit_corr * dec_tgt_mask) / tf.reduce_sum(dec_tgt_mask)) if FLAGS.edit_weight > 0: # monitor monitor_dict["edit_loss"] = edit_loss monitor_dict["accu_edit"] = edit_accu def get_class_acc(label_id): mask = tf.ones(tf.shape(edit_label), dtype=edit_label.dtype) * label_id mask = tf.equal(edit_label, mask) mask = tf.cast(mask, dtype=dec_tgt_mask.dtype) * dec_tgt_mask accu = tf.reduce_sum(edit_corr * mask) / tf.reduce_sum(mask) return accu monitor_dict["accu_orig"] = get_class_acc(0) if FLAGS.del_ratio > 0: monitor_dict["accu_del"] = get_class_acc(FLAGS.del_label) if FLAGS.ins_ratio > 0: monitor_dict["accu_ins"] = get_class_acc(FLAGS.ins_label) if FLAGS.rep_ratio > 0: monitor_dict["accu_rep"] = get_class_acc(FLAGS.rep_label) # accumulate total loss total_loss += FLAGS.edit_weight * edit_loss #### LM loss #### Only predict the target part tgt_out = output[:, src_len:] tf.logging.info("Output: %s, target output: %s", output.shape, tgt_out.shape) lm_loss, _ = model.lm_loss(hidden=tgt_out, target=dec_masked_tgt, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=shared_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=dec_mask_map, use_tpu=FLAGS.use_tpu) if lm_loss.dtype != tf.float32: lm_loss = tf.cast(lm_loss, tf.float32) dec_lm_tgt_mask = tf.cast(dec_lm_tgt_mask, lm_loss.dtype) lm_loss = (tf.reduce_sum(lm_loss * dec_lm_tgt_mask) / tf.reduce_sum(dec_lm_tgt_mask)) if FLAGS.lm_weight > 0: # monitor monitor_dict["lm_loss"] = lm_loss # accumulate total loss total_loss += FLAGS.lm_weight * lm_loss return total_loss, {}, monitor_dict
def electra_loss(features, labels, mems, n_token, is_training): """Electra Loss.""" del mems del labels initializer = _get_initializer() monitor_dict = {} #### Unpack input original_inp = features["inputs"] masked_inp = features["masked_input"] type_id = features["type_id"] target_mapping = features["target_mapping"] target = features["target"] gen_inp = masked_inp enc_type = type_id #### Shared input embedding (for generator) with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) gen_embed, shared_embed_table = inp_func(inputs=gen_inp, type_id=enc_type, return_embed_table=True) #### Generator TFM with tf.variable_scope("generator", reuse=tf.AUTO_REUSE): gen_tfm_func = _get_tfm_func(initializer, is_training, "pretrain", shrink=FLAGS.gen_shrink) gen_output, _ = gen_tfm_func(inputs=gen_embed, input_mask=None, perm_mask=None) gen_lm_loss, _, logits = model.lm_loss(hidden=gen_output, target=target, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=shared_embed_table, tie_weight=FLAGS.tie_weight, target_mapping=None, hidden_mapping=target_mapping, return_logits=True, use_tpu=FLAGS.use_tpu) assert FLAGS.sample_strategy != "token_span" if gen_lm_loss.dtype != tf.float32: gen_lm_loss = tf.cast(gen_lm_loss, tf.float32) gen_loss = tf.reduce_mean(gen_lm_loss) monitor_dict["gen_loss"] = gen_loss #### Sample from generator uniform = tf.random.uniform(minval=0, maxval=1, shape=logits.shape, dtype=tf.float32) gumbel = -tf.log(-tf.log(uniform + 1e-9) + 1e-9) samples = tf.argmax(logits + gumbel, -1) # map `num_predict` samples to full length samples = tf.einsum("bm,bml->bl", tf.cast(samples, tf.float32), tf.cast(target_mapping, tf.float32)) samples = tf.cast(samples, tf.int32) enc_inp = tf.where(tf.equal(gen_inp, FLAGS.mask_id), samples, gen_inp) is_diff = tf.cast(tf.not_equal(enc_inp, original_inp), tf.float32) monitor_dict["diff_percent"] = tf.reduce_mean(is_diff) edit_label = tf.cast(tf.equal(enc_inp, original_inp), tf.float32) #### Transformer Model with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_embed = inp_func(inputs=enc_inp, type_id=enc_type) tfm_func = _get_tfm_func(initializer, is_training, phase="pretrain") output, _ = tfm_func(inputs=inp_embed, input_mask=None, perm_mask=None) #### edit loss edit_loss, edit_corr = model.binary_loss(hidden=output, target=edit_label, d_model=FLAGS.d_model, initializer=initializer, target_mapping=None, hidden_mapping=None) edit_loss = tf.reduce_mean(edit_loss) edit_accu = tf.reduce_mean(edit_corr * is_diff) # monitor monitor_dict["edit_loss"] = edit_loss monitor_dict["edit_accu_fake"] = edit_accu # accumulate total loss edit_coeff = 50 total_loss = edit_coeff * edit_loss + gen_loss return total_loss, {}, monitor_dict
def joint_loss(features, labels, n_token, is_training): """Decoder only seq2seq.""" del labels initializer = _get_initializer() #### Unpack input source = features["source"] source_pos = features["source_position"] source_seg = features["source_segmentation"] target = features["target"] target_pos = features["target_position"] target_seg = features["target_segmentation"] # shapes bsz = tf.shape(source)[0] src_len = tf.shape(source)[1] tgt_len = tf.shape(target)[1] if FLAGS.use_bfloat16: tf_float = tf.bfloat16 else: tf_float = tf.float32 ##### format inputs inputs = tf.concat([source, target], axis=1) position = tf.concat([source_pos, target_pos], axis=1) src_type_id = features.get("source_type", tf.zeros([bsz, src_len], dtype=inputs.dtype)) tgt_type_id = features.get("target_type", tf.ones([bsz, tgt_len], dtype=inputs.dtype)) if FLAGS.double_type: tgt_type_id = tgt_type_id + FLAGS.n_token type_id = tf.concat([src_type_id, tgt_type_id], axis=1) ##### attention mask: note that `1` indicates CANNOT attend # src mask src_to_src = tf.not_equal(source_seg[:, :, None], source_seg[:, None, :]) src_to_tgt = tf.ones([bsz, src_len, tgt_len], dtype=src_to_src.dtype) src_mask = tf.concat([src_to_src, src_to_tgt], axis=2) # tgt mask tgt_to_src = tf.not_equal(target_seg[:, :, None], source_seg[:, None, :]) tgt_to_tgt = tf.not_equal(target_seg[:, :, None], target_seg[:, None, :]) causal_mask = tf.cast(causal_attn_mask(qlen=tgt_len), tgt_to_tgt.dtype) # If any one of them is `1` (indicating cannot attend), i.e. `logical_or`, # then the model should NOT attend tgt_to_tgt = tf.logical_or(tgt_to_tgt, causal_mask) tgt_mask = tf.concat([tgt_to_src, tgt_to_tgt], axis=2) # concat perm_mask = tf.concat([src_mask, tgt_mask], axis=1) perm_mask = tf.cast(perm_mask, tf_float) # padding non_pad_mask = tf.not_equal(target_seg, 0) all_eos = tf.constant(FLAGS.eos_id, shape=target.shape, dtype=target.dtype) # Replace all <pad> (/P) with <eos> (/S) # - target : /S a1 a2 a3 /S b1 b2 /S c1 c2 /P /P # - tmptgt : /S a1 a2 a3 /S b1 b2 /S c1 c2 /S /S tmptgt = tf.where(non_pad_mask, target, all_eos) # Shift the `tmptgt` to form the (next-step) prediction target # - target : \S a1 a2 a3 \S b1 b2 \S c1 c2 \P \P # - pred_tgt : a1 a2 a3 \S b1 b2 \S c1 c2 \S \S \S pred_tgt = tf.concat([tmptgt[:, 1:], tmptgt[:, :1]], axis=1) loss_mask = tf.cast(non_pad_mask, tf.float32) #### Transformer Model with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) input_embed, word_embed_table = inp_func(inputs=inputs, type_id=type_id, pos_seq=position, return_embed_table=True) tfm_func = _get_tfm_func(initializer, is_training, phase="pretrain") output, _ = tfm_func(inputs=input_embed, input_mask=None, perm_mask=perm_mask, pos_seq=position) #### Only predict the target part tgt_out = output[:, src_len:] tf.logging.info("Output: %s, target output: %s", output.shape, tgt_out.shape) lm_loss, nll_loss = model.lm_loss(hidden=tgt_out, target=pred_tgt, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=word_embed_table, tie_weight=FLAGS.tie_weight, label_smooth=FLAGS.label_smooth, use_tpu=FLAGS.use_tpu) if lm_loss.dtype != tf.float32: lm_loss = tf.cast(lm_loss, tf.float32) num_loss = tf.reduce_sum(loss_mask) total_loss = tf.reduce_sum(lm_loss * loss_mask) / num_loss nll = tf.reduce_sum(nll_loss * loss_mask) / num_loss # To be compatible with fairseq, convert to base 2 for logging monitor_dict = { "loss": total_loss / math.log(2), "nll": nll / math.log(2), "num_loss": num_loss, } return total_loss, monitor_dict
def mass_loss(features, labels, mems, n_token, is_training): """MASS pretraining loss.""" del labels del mems initializer = _get_initializer() # Type if FLAGS.use_bfloat16: tf_float = tf.bfloat16 else: tf_float = tf.float32 #### Unpack input target = features["target"] target_mask = features["target_mask"] target_mapping = features["target_mapping"] enc_inp = features["enc_inp"] enc_type = features["type_id"] enc_pos = features["enc_pos"] dec_inp = features["dec_inp"] dec_type = features["dec_type"] dec_pos = features["dec_pos"] dec_seg = features["dec_seg"] # shapes bsz = tf.shape(enc_inp)[0] enc_len = tf.shape(enc_inp)[1] dec_len = tf.shape(dec_inp)[1] ##### format inputs inputs = tf.concat([enc_inp, dec_inp], axis=1) position = tf.concat([enc_pos, dec_pos], axis=1) type_id = tf.concat([enc_type, dec_type], axis=1) ##### attention mask: note that `1` indicates CANNOT attend # enc mask enc_to_enc = tf.zeros([bsz, enc_len, enc_len], dtype=tf_float) enc_to_dec = tf.ones([bsz, enc_len, dec_len], dtype=tf_float) enc_mask = tf.concat([enc_to_enc, enc_to_dec], axis=2) # dec mask dec_to_enc = tf.zeros([bsz, dec_len, enc_len], dtype=tf.bool) dec_to_dec = tf.not_equal(dec_seg[:, :, None], dec_seg[:, None, :]) causal_mask = tf.cast(causal_attn_mask(qlen=dec_len), tf.bool) # If any one of them is `1` (indicating cannot attend), i.e. `logical_or`, # then the model should NOT attend dec_to_dec = tf.logical_or(dec_to_dec, causal_mask) dec_mask = tf.cast(tf.concat([dec_to_enc, dec_to_dec], axis=2), tf_float) # concat perm_mask = tf.concat([enc_mask, dec_mask], axis=1) perm_mask = tf.cast(perm_mask, tf_float) #### Transformer Model with tf.variable_scope("model", reuse=tf.AUTO_REUSE): inp_func = _get_inp_func(n_token, FLAGS.d_model, initializer, is_training) input_embed, word_embed_table = inp_func(inputs=inputs, type_id=type_id, pos_seq=position, return_embed_table=True) tfm_func = _get_tfm_func(initializer, is_training, phase="pretrain") output, _ = tfm_func(inputs=input_embed, input_mask=None, perm_mask=perm_mask) #### Only predict the target part enc_out = output[:, :enc_len] dec_out = output[:, enc_len:] tf.logging.info("Output: %s, enc output: %s, dec output: %s", output.shape, enc_out.shape, dec_out.shape) enc_loss, _ = model.lm_loss(hidden=enc_out, target=target, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=word_embed_table, tie_weight=FLAGS.tie_weight, hidden_mapping=target_mapping, use_tpu=FLAGS.use_tpu) dec_loss, _ = model.lm_loss(hidden=dec_out, target=target, n_token=n_token, d_model=FLAGS.d_model, initializer=initializer, lookup_table=word_embed_table, tie_weight=FLAGS.tie_weight, use_tpu=FLAGS.use_tpu) if dec_loss.dtype != tf.float32: dec_loss = tf.cast(dec_loss, tf.float32) if enc_loss.dtype != tf.float32: enc_loss = tf.cast(enc_loss, tf.float32) loss_mask = tf.cast(target_mask, tf.float32) num_loss = tf.reduce_sum(loss_mask) avg_dec_loss = tf.reduce_sum(dec_loss * loss_mask) / num_loss avg_enc_loss = tf.reduce_sum(enc_loss * loss_mask) / num_loss monitor_dict = {} total_loss = 0 if FLAGS.enc_weight > 0: total_loss += FLAGS.enc_weight * avg_enc_loss monitor_dict["loss_enc"] = avg_enc_loss if FLAGS.dec_weight > 0: total_loss += FLAGS.dec_weight * avg_dec_loss monitor_dict["loss_dec"] = avg_dec_loss monitor_dict["loss"] = total_loss return total_loss, {}, monitor_dict