Internal change

PiperOrigin-RevId: 333772508

official/nlp/modeling/models/bert_pretrainer.py

@@ -215,7 +215,7 @@ class BertPretrainerV2(tf.keras.Model):
     masked_lm_positions = tf.keras.layers.Input(
         shape=(None,), name='masked_lm_positions', dtype=tf.int32)
     inputs.append(masked_lm_positions)
-    outputs['lm_output'] = self.masked_lm(
+    outputs['mlm_logits'] = self.masked_lm(
         sequence_output, masked_positions=masked_lm_positions)
     for cls_head in self.classification_heads:
       outputs[cls_head.name] = cls_head(sequence_output)

official/nlp/modeling/models/bert_pretrainer_test.py

@@ -145,15 +145,16 @@ class BertPretrainerTest(keras_parameterized.TestCase):
     if has_encoder_outputs:
       self.assertSameElements(
           outputs.keys(),
-          ['sequence_output', 'pooled_output', 'lm_output', 'encoder_outputs'])
+          ['sequence_output', 'pooled_output', 'mlm_logits', 'encoder_outputs'])
       self.assertLen(outputs['encoder_outputs'], num_layers)
     else:
-      self.assertSameElements(outputs.keys(),
-                              ['sequence_output', 'pooled_output', 'lm_output'])
+      self.assertSameElements(
+          outputs.keys(), ['sequence_output', 'pooled_output', 'mlm_logits'])
 
     # Validate that the outputs are of the expected shape.
     expected_lm_shape = [None, num_token_predictions, vocab_size]
-    self.assertAllEqual(expected_lm_shape, outputs['lm_output'].shape.as_list())
+    self.assertAllEqual(expected_lm_shape,
+                        outputs['mlm_logits'].shape.as_list())
 
     expected_sequence_output_shape = [None, sequence_length, hidden_size]
     self.assertAllEqual(expected_sequence_output_shape,

official/nlp/tasks/masked_lm.py

@@ -67,7 +67,7 @@ class MaskedLMTask(base_task.Task):
       metrics = dict([(metric.name, metric) for metric in metrics])
       lm_prediction_losses = tf.keras.losses.sparse_categorical_crossentropy(
           labels['masked_lm_ids'],
-          tf.cast(model_outputs['lm_output'], tf.float32),
+          tf.cast(model_outputs['mlm_logits'], tf.float32),
           from_logits=True)
       lm_label_weights = labels['masked_lm_weights']
       lm_numerator_loss = tf.reduce_sum(lm_prediction_losses *
@@ -134,7 +134,7 @@ class MaskedLMTask(base_task.Task):
       metrics = dict([(metric.name, metric) for metric in metrics])
       if 'masked_lm_accuracy' in metrics:
         metrics['masked_lm_accuracy'].update_state(
-            labels['masked_lm_ids'], model_outputs['lm_output'],
+            labels['masked_lm_ids'], model_outputs['mlm_logits'],
             labels['masked_lm_weights'])
       if 'next_sentence_accuracy' in metrics:
         metrics['next_sentence_accuracy'].update_state(