Rename layers

luozhouyang · Nov 23, 2021 · 5910762 · 5910762
1 parent 16b0d43
commit 5910762
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/transformers_keras/question_answering/models.py b/transformers_keras/question_answering/models.py
@@ -45,9 +45,9 @@ def __init__(
 
         sequence_output, _, _, _ = bert_model(input_ids, segment_ids, attention_mask)
         logits = tf.keras.layers.Dense(num_labels, name="dense")(sequence_output)
-        head_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="head")(logits)
-        tail_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="tail")(logits)
-        super().__init__(inputs=[input_ids, segment_ids, attention_mask], outputs=[head_logits, tail_logits], **kwargs)
+        start_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="start")(logits)
+        end_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="end")(logits)
+        super().__init__(inputs=[input_ids, segment_ids, attention_mask], outputs=[start_logits, end_logits], **kwargs)
 
 
 class AlbertForQuestionAnswering(AlbertPretrainedModel):
@@ -98,9 +98,9 @@ def __init__(
 
         sequence_output, _, _, _ = albert_model(input_ids, segment_ids, attention_mask)
         logits = tf.keras.layers.Dense(num_labels, name="dense")(sequence_output)
-        head_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="head")(logits)
-        tail_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="tail")(logits)
-        super().__init__(inputs=[input_ids, segment_ids, attention_mask], outputs=[head_logits, tail_logits], **kwargs)
+        start_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="start")(logits)
+        end_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="end")(logits)
+        super().__init__(inputs=[input_ids, segment_ids, attention_mask], outputs=[start_logits, end_logits], **kwargs)
 
 
 class BertForQuestionAnsweringX(BertPretrainedModel):
@@ -146,12 +146,12 @@ def __init__(
 
         sequence_output, pooled_output, _, _ = bert_model(input_ids, segment_ids, attention_mask)
         logits = tf.keras.layers.Dense(num_labels, name="dense")(sequence_output)
-        head_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="head")(logits)
-        tail_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="tail")(logits)
+        start_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="start")(logits)
+        end_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="end")(logits)
         class_logits = tf.keras.layers.Dense(num_classes, name="class")(pooled_output)
 
         super().__init__(
-            inputs=[input_ids, segment_ids, attention_mask], outputs=[head_logits, tail_logits, class_logits], **kwargs
+            inputs=[input_ids, segment_ids, attention_mask], outputs=[start_logits, end_logits, class_logits], **kwargs
         )
 
 
@@ -204,9 +204,9 @@ def __init__(
 
         sequence_output, pooled_output, _, _ = albert_model(input_ids, segment_ids, attention_mask)
         logits = tf.keras.layers.Dense(num_labels, name="dense")(sequence_output)
-        head_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="head")(logits)
-        tail_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="tail")(logits)
+        start_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 0], name="start")(logits)
+        end_logits = tf.keras.layers.Lambda(lambda x: x[:, :, 1], name="end")(logits)
         class_logits = tf.keras.layers.Dense(num_classes, name="class")(pooled_output)
         super().__init__(
-            inputs=[input_ids, segment_ids, attention_mask], outputs=[head_logits, tail_logits, class_logits], **kwargs
+            inputs=[input_ids, segment_ids, attention_mask], outputs=[start_logits, end_logits, class_logits], **kwargs
         )