Internal change

PiperOrigin-RevId: 398593113

Author: tensorflower-gardener

official/core/input_reader.py

@@ -204,6 +204,7 @@ class EncoderConfig(hyperparams.OneOfConfig):
   bigbird: BigBirdEncoderConfig = BigBirdEncoderConfig()
   kernel: KernelEncoderConfig = KernelEncoderConfig()
   mobilebert: MobileBertEncoderConfig = MobileBertEncoderConfig()
+  teams: BertEncoderConfig = BertEncoderConfig()
   xlnet: XLNetEncoderConfig = XLNetEncoderConfig()
 
 
@@ -436,6 +437,40 @@ def build_encoder(config: EncoderConfig,
         initializer=tf.keras.initializers.RandomNormal(
             stddev=encoder_cfg.initializer_range))
 
+  if encoder_type == "teams":
+    embedding_cfg = dict(
+        vocab_size=encoder_cfg.vocab_size,
+        type_vocab_size=encoder_cfg.type_vocab_size,
+        hidden_size=encoder_cfg.hidden_size,
+        embedding_width=encoder_cfg.embedding_size,
+        max_seq_length=encoder_cfg.max_position_embeddings,
+        initializer=tf.keras.initializers.TruncatedNormal(
+            stddev=encoder_cfg.initializer_range),
+        dropout_rate=encoder_cfg.dropout_rate,
+    )
+    embedding_network = networks.PackedSequenceEmbedding(**embedding_cfg)
+    hidden_cfg = dict(
+        num_attention_heads=encoder_cfg.num_attention_heads,
+        intermediate_size=encoder_cfg.intermediate_size,
+        intermediate_activation=tf_utils.get_activation(
+            encoder_cfg.hidden_activation),
+        dropout_rate=encoder_cfg.dropout_rate,
+        attention_dropout_rate=encoder_cfg.attention_dropout_rate,
+        kernel_initializer=tf.keras.initializers.TruncatedNormal(
+            stddev=encoder_cfg.initializer_range),
+    )
+    kwargs = dict(
+        embedding_cfg=embedding_cfg,
+        embedding_cls=embedding_network,
+        hidden_cfg=hidden_cfg,
+        num_hidden_instances=encoder_cfg.num_layers,
+        pooled_output_dim=encoder_cfg.hidden_size,
+        pooler_layer_initializer=tf.keras.initializers.TruncatedNormal(
+            stddev=encoder_cfg.initializer_range),
+        return_all_layer_outputs=encoder_cfg.return_all_encoder_outputs,
+        dict_outputs=True)
+    return networks.EncoderScaffold(**kwargs)
+
   # Uses the default BERTEncoder configuration schema to create the encoder.
   # If it does not match, please add a switch branch by the encoder type.
   return networks.BertEncoder(

official/nlp/configs/encoders.py

@@ -61,7 +61,6 @@ def bert_sentence_prediction() -> cfg.ExperimentConfig:
           'task.train_data.is_training != None',
           'task.validation_data.is_training != None'
       ])
-  config.task.model.encoder.type = 'bert'
   return config
 
 
@@ -98,7 +97,6 @@ def bert_squad() -> cfg.ExperimentConfig:
           'task.train_data.is_training != None',
           'task.validation_data.is_training != None'
       ])
-  config.task.model.encoder.type = 'bert'
   return config
 
 

official/nlp/configs/finetuning_experiments.py

@@ -14,23 +14,24 @@
 
 """Funnel Transformer network."""
 # pylint: disable=g-classes-have-attributes
-from typing import Union, Collection
+from typing import Union, Sequence
 from absl import logging
 import tensorflow as tf
 
 from official.nlp import keras_nlp
 
 
-def _pool_and_concat(data, unpool_length: int, stride: int,
-                     axes: Union[Collection[int], int]):
+def _pool_and_concat(data, unpool_length: int, strides: Union[Sequence[int],
+                                                              int],
+                     axes: Union[Sequence[int], int]):
   """Pools the data along a given axis with stride.
 
   It also skips first unpool_length elements.
 
   Args:
     data: Tensor to be pooled.
     unpool_length: Leading elements to be skipped.
-    stride: Stride for the given axis.
+    strides: Strides for the given axes.
     axes: Axes to pool the Tensor.
 
   Returns:
@@ -39,8 +40,13 @@ def _pool_and_concat(data, unpool_length: int, stride: int,
   # Wraps the axes as a list.
   if isinstance(axes, int):
     axes = [axes]
+  if isinstance(strides, int):
+    strides = [strides] * len(axes)
+  else:
+    if len(strides) != len(axes):
+      raise ValueError('The lengths of strides and axes need to match.')
 
-  for axis in axes:
+  for axis, stride in zip(axes, strides):
     # Skips first `unpool_length` tokens.
     unpool_tensor_shape = [slice(None)] * axis + [slice(None, unpool_length)]
     unpool_tensor = data[unpool_tensor_shape]
@@ -80,7 +86,9 @@ class FunnelTransformerEncoder(tf.keras.layers.Layer):
       dropout.
     attention_dropout: The dropout rate to use for the attention layers within
       the transformer layers.
-    pool_stride: Pooling stride to compress the sequence length.
+    pool_stride: An int or a list of ints. Pooling stride(s) to compress the
+      sequence length. If set to int, each layer will have the same stride size.
+      If set to list, the number of elements needs to match num_layers.
     unpool_length: Leading n tokens to be skipped from pooling.
     initializer: The initialzer to use for all weights in this encoder.
     output_range: The sequence output range, [0, output_range), by slicing the
@@ -185,12 +193,23 @@ def __init__(
         activation='tanh',
         kernel_initializer=initializer,
         name='pooler_transform')
-    self._att_input_pool_layer = tf.keras.layers.MaxPooling1D(
-        pool_size=pool_stride,
-        strides=pool_stride,
-        padding='same',
-        name='att_input_pool_layer')
-    self._pool_stride = pool_stride
+    if isinstance(pool_stride, int):
+      # TODO(b/197133196): Pooling layer can be shared.
+      pool_strides = [pool_stride] * num_layers
+    else:
+      if len(pool_stride) != num_layers:
+        raise ValueError('Lengths of pool_stride and num_layers are not equal.')
+      pool_strides = pool_stride
+    self._att_input_pool_layers = []
+    for layer_pool_stride in pool_strides:
+      att_input_pool_layer = tf.keras.layers.MaxPooling1D(
+          pool_size=layer_pool_stride,
+          strides=layer_pool_stride,
+          padding='same',
+          name='att_input_pool_layer')
+      self._att_input_pool_layers.append(att_input_pool_layer)
+
+    self._pool_strides = pool_strides  # This is a list here.
     self._unpool_length = unpool_length
 
     self._config = {
@@ -250,23 +269,25 @@ def call(self, inputs):
     attention_mask = _pool_and_concat(
         attention_mask,
         unpool_length=self._unpool_length,
-        stride=self._pool_stride,
+        strides=self._pool_strides[0],
         axes=[1])
-    for layer in self._transformer_layers:
+    for i, layer in enumerate(self._transformer_layers):
       # Pools layer for compressing the query length.
-      pooled_inputs = self._att_input_pool_layer(x[:, self._unpool_length:, :])
+      pooled_inputs = self._att_input_pool_layers[i](
+          x[:, self._unpool_length:, :])
       query_inputs = tf.concat(
           values=(tf.cast(
               x[:, :self._unpool_length, :],
               dtype=pooled_inputs.dtype), pooled_inputs),
           axis=1)
       x = layer([query_inputs, x, attention_mask])
       # Pools the corresponding attention_mask.
-      attention_mask = _pool_and_concat(
-          attention_mask,
-          unpool_length=self._unpool_length,
-          stride=self._pool_stride,
-          axes=[1, 2])
+      if i < len(self._transformer_layers) - 1:
+        attention_mask = _pool_and_concat(
+            attention_mask,
+            unpool_length=self._unpool_length,
+            strides=[self._pool_strides[i+1], self._pool_strides[i]],
+            axes=[1, 2])
       encoder_outputs.append(x)
 
     last_encoder_output = encoder_outputs[-1]

official/nlp/modeling/networks/funnel_transformer.py

@@ -80,8 +80,24 @@ def test_network_creation(self, policy, pooled_dtype):
     self.assertAllEqual(tf.float32, data.dtype)
     self.assertAllEqual(pooled_dtype, pooled.dtype)
 
+  def test_invalid_stride_and_num_layers(self):
+    hidden_size = 32
+    num_layers = 3
+    pool_stride = [2, 2]
+    unpool_length = 1
+    with self.assertRaisesRegex(ValueError,
+                                "pool_stride and num_layers are not equal"):
+      _ = funnel_transformer.FunnelTransformerEncoder(
+          vocab_size=100,
+          hidden_size=hidden_size,
+          num_attention_heads=2,
+          num_layers=num_layers,
+          pool_stride=pool_stride,
+          unpool_length=unpool_length)
+
   @parameterized.named_parameters(
       ("no_stride_no_unpool", 1, 0),
+      ("stride_list_with_unpool", [2, 3, 4], 1),
       ("large_stride_with_unpool", 3, 1),
       ("large_stride_with_large_unpool", 5, 10),
       ("no_stride_with_unpool", 1, 1),
@@ -110,11 +126,12 @@ def test_all_encoder_outputs_network_creation(self, pool_stride,
     expected_data_shape = [None, sequence_length, hidden_size]
     expected_pooled_shape = [None, hidden_size]
     self.assertLen(all_encoder_outputs, num_layers)
-    for data in all_encoder_outputs:
-      expected_data_shape[1] = unpool_length + (expected_data_shape[1] +
-                                                pool_stride - 1 -
-                                                unpool_length) // pool_stride
-      print("shapes:", expected_data_shape, data.shape.as_list())
+    if isinstance(pool_stride, int):
+      pool_stride = [pool_stride] * num_layers
+    for layer_pool_stride, data in zip(pool_stride, all_encoder_outputs):
+      expected_data_shape[1] = unpool_length + (
+          expected_data_shape[1] + layer_pool_stride - 1 -
+          unpool_length) // layer_pool_stride
       self.assertAllEqual(expected_data_shape, data.shape.as_list())
     self.assertAllEqual(expected_pooled_shape, pooled.shape.as_list())
 

official/nlp/modeling/networks/funnel_transformer_test.py

@@ -62,6 +62,8 @@ def __init__(self,
                pack_multiple_sequences=False,
                **kwargs):
     initializer = tf.keras.initializers.get(initializer)
+    if embedding_width is None:
+      embedding_width = hidden_size
     config_dict = {
         'vocab_size': vocab_size,
         'type_vocab_size': type_vocab_size,

official/nlp/modeling/networks/packed_sequence_embedding.py

@@ -1,7 +1,7 @@
 task:
   model:
     encoder:
-      bert:
+      teams:
         attention_dropout_rate: 0.1
         dropout_rate: 0.1
         embedding_size: 768
@@ -14,3 +14,4 @@ task:
         num_layers: 12
         type_vocab_size: 2
         vocab_size: 30522
+      type: teams

official/nlp/projects/teams/experiments/teams_en_uncased_base.yaml

@@ -1,7 +1,7 @@
 task:
   model:
     encoder:
-      bert:
+      teams:
         attention_dropout_rate: 0.1
         dropout_rate: 0.1
         embedding_size: 128
@@ -14,3 +14,4 @@ task:
         num_layers: 12
         type_vocab_size: 2
         vocab_size: 30522
+      type: teams

official/nlp/projects/teams/experiments/teams_en_uncased_small.yaml

@@ -64,9 +64,6 @@ def get_encoder(bert_config,
   Returns:
     A encoder object.
   """
-  # embedding_size is required for PackedSequenceEmbedding.
-  if bert_config.embedding_size is None:
-    bert_config.embedding_size = bert_config.hidden_size
   embedding_cfg = dict(
       vocab_size=bert_config.vocab_size,
       type_vocab_size=bert_config.type_vocab_size,

official/nlp/projects/teams/teams.py

@@ -21,6 +21,7 @@ pyyaml>=5.1
 opencv-python-headless
 Pillow
 pycocotools
+waymo-open-dataset-tf-2-6-0
 # NLP related dependencies
 seqeval
 sentencepiece