encoders.py 20 KB
Newer Older
Frederick Liu's avatar
Frederick Liu committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Frederick Liu's avatar
Frederick Liu committed
14

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
15
16
"""Transformer Encoders.

Hongkun Yu's avatar
Hongkun Yu committed
17
Includes configurations and factory methods.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
18
"""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
19
from typing import Optional
Hongkun Yu's avatar
Hongkun Yu committed
20

21
import dataclasses
Hongkun Yu's avatar
Hongkun Yu committed
22
import gin
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
23
import tensorflow as tf
24

Hongkun Yu's avatar
Hongkun Yu committed
25
from official.modeling import hyperparams
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
26
from official.modeling import tf_utils
Frederick Liu's avatar
Frederick Liu committed
27
from official.nlp.modeling import layers
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
28
from official.nlp.modeling import networks
Hongkun Yu's avatar
Hongkun Yu committed
29
from official.nlp.projects.bigbird import encoder as bigbird_encoder
30
31
32


@dataclasses.dataclass
Hongkun Yu's avatar
Hongkun Yu committed
33
class BertEncoderConfig(hyperparams.Config):
34
35
36
37
38
39
  """BERT encoder configuration."""
  vocab_size: int = 30522
  hidden_size: int = 768
  num_layers: int = 12
  num_attention_heads: int = 12
  hidden_activation: str = "gelu"
Chen Chen's avatar
Chen Chen committed
40
  intermediate_size: int = 3072
41
42
43
44
45
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
  max_position_embeddings: int = 512
  type_vocab_size: int = 2
  initializer_range: float = 0.02
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
46
  embedding_size: Optional[int] = None
Frederick Liu's avatar
Frederick Liu committed
47
  output_range: Optional[int] = None
Chen Chen's avatar
Chen Chen committed
48
  return_all_encoder_outputs: bool = False
49
50
  # Pre/Post-LN Transformer
  norm_first: bool = False
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
51
52


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
53
54
55
56
57
58
59
60
61
62
63
64
@dataclasses.dataclass
class MobileBertEncoderConfig(hyperparams.Config):
  """MobileBERT encoder configuration.

  Attributes:
    word_vocab_size: number of words in the vocabulary.
    word_embed_size: word embedding size.
    type_vocab_size: number of word types.
    max_sequence_length: maximum length of input sequence.
    num_blocks: number of transformer block in the encoder model.
    hidden_size: the hidden size for the transformer block.
    num_attention_heads: number of attention heads in the transformer block.
Hongkun Yu's avatar
Hongkun Yu committed
65
66
    intermediate_size: the size of the "intermediate" (a.k.a., feed forward)
      layer.
Chen Chen's avatar
Chen Chen committed
67
    hidden_activation: the non-linear activation function to apply to the
Hongkun Yu's avatar
Hongkun Yu committed
68
      output of the intermediate/feed-forward layer.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
69
70
71
72
73
    hidden_dropout_prob: dropout probability for the hidden layers.
    attention_probs_dropout_prob: dropout probability of the attention
      probabilities.
    intra_bottleneck_size: the size of bottleneck.
    initializer_range: The stddev of the truncated_normal_initializer for
Hongkun Yu's avatar
Hongkun Yu committed
74
      initializing all weight matrices.
Chen Chen's avatar
Chen Chen committed
75
76
77
    use_bottleneck_attention: Use attention inputs from the bottleneck
      transformation. If true, the following `key_query_shared_bottleneck`
      will be ignored.
Hongkun Yu's avatar
Hongkun Yu committed
78
79
    key_query_shared_bottleneck: whether to share linear transformation for keys
      and queries.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    num_feedforward_networks: number of stacked feed-forward networks.
    normalization_type: the type of normalization_type, only 'no_norm' and
      'layer_norm' are supported. 'no_norm' represents the element-wise linear
      transformation for the student model, as suggested by the original
      MobileBERT paper. 'layer_norm' is used for the teacher model.
    classifier_activation: if using the tanh activation for the final
      representation of the [CLS] token in fine-tuning.
  """
  word_vocab_size: int = 30522
  word_embed_size: int = 128
  type_vocab_size: int = 2
  max_sequence_length: int = 512
  num_blocks: int = 24
  hidden_size: int = 512
  num_attention_heads: int = 4
  intermediate_size: int = 4096
Chen Chen's avatar
Chen Chen committed
96
  hidden_activation: str = "gelu"
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
97
98
99
100
  hidden_dropout_prob: float = 0.1
  attention_probs_dropout_prob: float = 0.1
  intra_bottleneck_size: int = 1024
  initializer_range: float = 0.02
Chen Chen's avatar
Chen Chen committed
101
  use_bottleneck_attention: bool = False
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
102
103
104
105
  key_query_shared_bottleneck: bool = False
  num_feedforward_networks: int = 1
  normalization_type: str = "layer_norm"
  classifier_activation: bool = True
Chen Chen's avatar
Chen Chen committed
106
  input_mask_dtype: str = "int32"
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
107
108


Chen Chen's avatar
Chen Chen committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
@dataclasses.dataclass
class AlbertEncoderConfig(hyperparams.Config):
  """ALBERT encoder configuration."""
  vocab_size: int = 30000
  embedding_width: int = 128
  hidden_size: int = 768
  num_layers: int = 12
  num_attention_heads: int = 12
  hidden_activation: str = "gelu"
  intermediate_size: int = 3072
  dropout_rate: float = 0.0
  attention_dropout_rate: float = 0.0
  max_position_embeddings: int = 512
  type_vocab_size: int = 2
  initializer_range: float = 0.02


Hongkun Yu's avatar
Hongkun Yu committed
126
127
128
129
130
131
132
133
134
135
136
@dataclasses.dataclass
class BigBirdEncoderConfig(hyperparams.Config):
  """BigBird encoder configuration."""
  vocab_size: int = 50358
  hidden_size: int = 768
  num_layers: int = 12
  num_attention_heads: int = 12
  hidden_activation: str = "gelu"
  intermediate_size: int = 3072
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
137
138
  # Pre/Post-LN Transformer
  norm_first: bool = False
Hongkun Yu's avatar
Hongkun Yu committed
139
140
141
142
143
  max_position_embeddings: int = 4096
  num_rand_blocks: int = 3
  block_size: int = 64
  type_vocab_size: int = 16
  initializer_range: float = 0.02
144
  embedding_width: Optional[int] = None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
145
  use_gradient_checkpointing: bool = False
Hongkun Yu's avatar
Hongkun Yu committed
146
147


Frederick Liu's avatar
Frederick Liu committed
148
149
150
151
152
153
154
155
156
157
158
@dataclasses.dataclass
class KernelEncoderConfig(hyperparams.Config):
  """Linear encoder configuration."""
  vocab_size: int = 30522
  hidden_size: int = 768
  num_layers: int = 12
  num_attention_heads: int = 12
  hidden_activation: str = "gelu"
  intermediate_size: int = 3072
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
159
160
  # Pre/Post-LN Transformer
  norm_first: bool = False
Frederick Liu's avatar
Frederick Liu committed
161
162
163
164
165
166
167
168
169
  max_position_embeddings: int = 512
  type_vocab_size: int = 2
  initializer_range: float = 0.02
  embedding_size: Optional[int] = None
  feature_transform: str = "exp"
  num_random_features: int = 256
  redraw: bool = False
  is_short_seq: bool = False
  begin_kernel: int = 0
Frederick Liu's avatar
Frederick Liu committed
170
  scale: Optional[float] = None
Frederick Liu's avatar
Frederick Liu committed
171
172


Allen Wang's avatar
Allen Wang committed
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
@dataclasses.dataclass
class XLNetEncoderConfig(hyperparams.Config):
  """XLNet encoder configuration."""
  vocab_size: int = 32000
  num_layers: int = 24
  hidden_size: int = 1024
  num_attention_heads: int = 16
  head_size: int = 64
  inner_size: int = 4096
  inner_activation: str = "gelu"
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
  attention_type: str = "bi"
  bi_data: bool = False
  tie_attention_biases: bool = False
  memory_length: int = 0
  same_length: bool = False
  clamp_length: int = -1
  reuse_length: int = 0
  use_cls_mask: bool = False
  embedding_width: int = 1024
  initializer_range: float = 0.02
  two_stream: bool = False


Hongkun Yu's avatar
Hongkun Yu committed
198
199
200
201
@dataclasses.dataclass
class EncoderConfig(hyperparams.OneOfConfig):
  """Encoder configuration."""
  type: Optional[str] = "bert"
Chen Chen's avatar
Chen Chen committed
202
  albert: AlbertEncoderConfig = AlbertEncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
203
  bert: BertEncoderConfig = BertEncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
204
  bigbird: BigBirdEncoderConfig = BigBirdEncoderConfig()
Frederick Liu's avatar
Frederick Liu committed
205
  kernel: KernelEncoderConfig = KernelEncoderConfig()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
206
  mobilebert: MobileBertEncoderConfig = MobileBertEncoderConfig()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
207
  teams: BertEncoderConfig = BertEncoderConfig()
Allen Wang's avatar
Allen Wang committed
208
  xlnet: XLNetEncoderConfig = XLNetEncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
209
210
211


@gin.configurable
Hongkun Yu's avatar
Hongkun Yu committed
212
213
214
215
def build_encoder(config: EncoderConfig,
                  embedding_layer: Optional[tf.keras.layers.Layer] = None,
                  encoder_cls=None,
                  bypass_config: bool = False):
Hongkun Yu's avatar
Hongkun Yu committed
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
  """Instantiate a Transformer encoder network from EncoderConfig.

  Args:
    config: the one-of encoder config, which provides encoder parameters of a
      chosen encoder.
    embedding_layer: an external embedding layer passed to the encoder.
    encoder_cls: an external encoder cls not included in the supported encoders,
      usually used by gin.configurable.
    bypass_config: whether to ignore config instance to create the object with
      `encoder_cls`.

  Returns:
    An encoder instance.
  """
  if bypass_config:
    return encoder_cls()
Frederick Liu's avatar
Frederick Liu committed
232
233
234
  encoder_type = config.type
  encoder_cfg = config.get()
  if encoder_cls and encoder_cls.__name__ == "EncoderScaffold":
Hongkun Yu's avatar
Hongkun Yu committed
235
    embedding_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
236
237
238
239
        vocab_size=encoder_cfg.vocab_size,
        type_vocab_size=encoder_cfg.type_vocab_size,
        hidden_size=encoder_cfg.hidden_size,
        max_seq_length=encoder_cfg.max_position_embeddings,
Hongkun Yu's avatar
Hongkun Yu committed
240
        initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
241
242
            stddev=encoder_cfg.initializer_range),
        dropout_rate=encoder_cfg.dropout_rate,
Hongkun Yu's avatar
Hongkun Yu committed
243
244
    )
    hidden_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
245
246
        num_attention_heads=encoder_cfg.num_attention_heads,
        intermediate_size=encoder_cfg.intermediate_size,
Hongkun Yu's avatar
Hongkun Yu committed
247
        intermediate_activation=tf_utils.get_activation(
Hongkun Yu's avatar
Hongkun Yu committed
248
249
250
            encoder_cfg.hidden_activation),
        dropout_rate=encoder_cfg.dropout_rate,
        attention_dropout_rate=encoder_cfg.attention_dropout_rate,
Hongkun Yu's avatar
Hongkun Yu committed
251
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
252
            stddev=encoder_cfg.initializer_range),
Hongkun Yu's avatar
Hongkun Yu committed
253
254
255
256
    )
    kwargs = dict(
        embedding_cfg=embedding_cfg,
        hidden_cfg=hidden_cfg,
Hongkun Yu's avatar
Hongkun Yu committed
257
258
        num_hidden_instances=encoder_cfg.num_layers,
        pooled_output_dim=encoder_cfg.hidden_size,
Hongkun Yu's avatar
Hongkun Yu committed
259
        pooler_layer_initializer=tf.keras.initializers.TruncatedNormal(
Chen Chen's avatar
Chen Chen committed
260
            stddev=encoder_cfg.initializer_range),
261
262
        return_all_layer_outputs=encoder_cfg.return_all_encoder_outputs,
        dict_outputs=True)
Hongkun Yu's avatar
Hongkun Yu committed
263
264
    return encoder_cls(**kwargs)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
265
  if encoder_type == "mobilebert":
Frederick Liu's avatar
Frederick Liu committed
266
    return networks.MobileBERTEncoder(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
267
268
269
270
271
272
273
274
        word_vocab_size=encoder_cfg.word_vocab_size,
        word_embed_size=encoder_cfg.word_embed_size,
        type_vocab_size=encoder_cfg.type_vocab_size,
        max_sequence_length=encoder_cfg.max_sequence_length,
        num_blocks=encoder_cfg.num_blocks,
        hidden_size=encoder_cfg.hidden_size,
        num_attention_heads=encoder_cfg.num_attention_heads,
        intermediate_size=encoder_cfg.intermediate_size,
Chen Chen's avatar
Chen Chen committed
275
        intermediate_act_fn=encoder_cfg.hidden_activation,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
276
277
278
        hidden_dropout_prob=encoder_cfg.hidden_dropout_prob,
        attention_probs_dropout_prob=encoder_cfg.attention_probs_dropout_prob,
        intra_bottleneck_size=encoder_cfg.intra_bottleneck_size,
Chen Chen's avatar
Chen Chen committed
279
        initializer_range=encoder_cfg.initializer_range,
Chen Chen's avatar
Chen Chen committed
280
        use_bottleneck_attention=encoder_cfg.use_bottleneck_attention,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
281
282
283
        key_query_shared_bottleneck=encoder_cfg.key_query_shared_bottleneck,
        num_feedforward_networks=encoder_cfg.num_feedforward_networks,
        normalization_type=encoder_cfg.normalization_type,
Chen Chen's avatar
Chen Chen committed
284
285
        classifier_activation=encoder_cfg.classifier_activation,
        input_mask_dtype=encoder_cfg.input_mask_dtype)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
286

Chen Chen's avatar
Chen Chen committed
287
  if encoder_type == "albert":
Frederick Liu's avatar
Frederick Liu committed
288
    return networks.AlbertEncoder(
Chen Chen's avatar
Chen Chen committed
289
290
291
292
293
294
295
296
297
298
299
300
        vocab_size=encoder_cfg.vocab_size,
        embedding_width=encoder_cfg.embedding_width,
        hidden_size=encoder_cfg.hidden_size,
        num_layers=encoder_cfg.num_layers,
        num_attention_heads=encoder_cfg.num_attention_heads,
        max_sequence_length=encoder_cfg.max_position_embeddings,
        type_vocab_size=encoder_cfg.type_vocab_size,
        intermediate_size=encoder_cfg.intermediate_size,
        activation=tf_utils.get_activation(encoder_cfg.hidden_activation),
        dropout_rate=encoder_cfg.dropout_rate,
        attention_dropout_rate=encoder_cfg.attention_dropout_rate,
        initializer=tf.keras.initializers.TruncatedNormal(
301
302
            stddev=encoder_cfg.initializer_range),
        dict_outputs=True)
Chen Chen's avatar
Chen Chen committed
303

Hongkun Yu's avatar
Hongkun Yu committed
304
  if encoder_type == "bigbird":
Hongkun Yu's avatar
Hongkun Yu committed
305
306
    # TODO(frederickliu): Support use_gradient_checkpointing and update
    # experiments to use the EncoderScaffold only.
Frederick Liu's avatar
Frederick Liu committed
307
    if encoder_cfg.use_gradient_checkpointing:
Hongkun Yu's avatar
Hongkun Yu committed
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
      return bigbird_encoder.BigBirdEncoder(
          vocab_size=encoder_cfg.vocab_size,
          hidden_size=encoder_cfg.hidden_size,
          num_layers=encoder_cfg.num_layers,
          num_attention_heads=encoder_cfg.num_attention_heads,
          intermediate_size=encoder_cfg.intermediate_size,
          activation=tf_utils.get_activation(encoder_cfg.hidden_activation),
          dropout_rate=encoder_cfg.dropout_rate,
          attention_dropout_rate=encoder_cfg.attention_dropout_rate,
          num_rand_blocks=encoder_cfg.num_rand_blocks,
          block_size=encoder_cfg.block_size,
          max_position_embeddings=encoder_cfg.max_position_embeddings,
          type_vocab_size=encoder_cfg.type_vocab_size,
          initializer=tf.keras.initializers.TruncatedNormal(
              stddev=encoder_cfg.initializer_range),
          embedding_width=encoder_cfg.embedding_width,
          use_gradient_checkpointing=encoder_cfg.use_gradient_checkpointing)
Frederick Liu's avatar
Frederick Liu committed
325
    embedding_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
326
        vocab_size=encoder_cfg.vocab_size,
Frederick Liu's avatar
Frederick Liu committed
327
        type_vocab_size=encoder_cfg.type_vocab_size,
Hongkun Yu's avatar
Hongkun Yu committed
328
        hidden_size=encoder_cfg.hidden_size,
Frederick Liu's avatar
Frederick Liu committed
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
        max_seq_length=encoder_cfg.max_position_embeddings,
        initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        dropout_rate=encoder_cfg.dropout_rate)
    attention_cfg = dict(
        num_heads=encoder_cfg.num_attention_heads,
        key_dim=int(encoder_cfg.hidden_size // encoder_cfg.num_attention_heads),
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        max_rand_mask_length=encoder_cfg.max_position_embeddings,
        num_rand_blocks=encoder_cfg.num_rand_blocks,
        from_block_size=encoder_cfg.block_size,
        to_block_size=encoder_cfg.block_size,
        )
    hidden_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
344
345
        num_attention_heads=encoder_cfg.num_attention_heads,
        intermediate_size=encoder_cfg.intermediate_size,
Frederick Liu's avatar
Frederick Liu committed
346
347
        intermediate_activation=tf_utils.get_activation(
            encoder_cfg.hidden_activation),
Hongkun Yu's avatar
Hongkun Yu committed
348
349
        dropout_rate=encoder_cfg.dropout_rate,
        attention_dropout_rate=encoder_cfg.attention_dropout_rate,
350
        norm_first=encoder_cfg.norm_first,
Frederick Liu's avatar
Frederick Liu committed
351
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
352
            stddev=encoder_cfg.initializer_range),
353
        attention_cls=layers.BigBirdAttention,
Frederick Liu's avatar
Frederick Liu committed
354
355
356
357
358
359
        attention_cfg=attention_cfg)
    kwargs = dict(
        embedding_cfg=embedding_cfg,
        hidden_cls=layers.TransformerScaffold,
        hidden_cfg=hidden_cfg,
        num_hidden_instances=encoder_cfg.num_layers,
360
        mask_cls=layers.BigBirdMasks,
Frederick Liu's avatar
Frederick Liu committed
361
362
363
364
365
366
367
368
        mask_cfg=dict(block_size=encoder_cfg.block_size),
        pooled_output_dim=encoder_cfg.hidden_size,
        pooler_layer_initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        return_all_layer_outputs=False,
        dict_outputs=True,
        layer_idx_as_attention_seed=True)
    return networks.EncoderScaffold(**kwargs)
Hongkun Yu's avatar
Hongkun Yu committed
369

Frederick Liu's avatar
Frederick Liu committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
  if encoder_type == "kernel":
    embedding_cfg = dict(
        vocab_size=encoder_cfg.vocab_size,
        type_vocab_size=encoder_cfg.type_vocab_size,
        hidden_size=encoder_cfg.hidden_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)
    attention_cfg = dict(
        num_heads=encoder_cfg.num_attention_heads,
        key_dim=int(encoder_cfg.hidden_size // encoder_cfg.num_attention_heads),
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        feature_transform=encoder_cfg.feature_transform,
        num_random_features=encoder_cfg.num_random_features,
        redraw=encoder_cfg.redraw,
        is_short_seq=encoder_cfg.is_short_seq,
        begin_kernel=encoder_cfg.begin_kernel,
Frederick Liu's avatar
Frederick Liu committed
389
        scale=encoder_cfg.scale,
Frederick Liu's avatar
Frederick Liu committed
390
391
392
393
394
395
396
397
        )
    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,
398
        norm_first=encoder_cfg.norm_first,
Frederick Liu's avatar
Frederick Liu committed
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        attention_cls=layers.KernelAttention,
        attention_cfg=attention_cfg)
    kwargs = dict(
        embedding_cfg=embedding_cfg,
        hidden_cls=layers.TransformerScaffold,
        hidden_cfg=hidden_cfg,
        num_hidden_instances=encoder_cfg.num_layers,
        mask_cls=layers.KernelMask,
        pooled_output_dim=encoder_cfg.hidden_size,
        pooler_layer_initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        return_all_layer_outputs=False,
        dict_outputs=True,
        layer_idx_as_attention_seed=True)
    return networks.EncoderScaffold(**kwargs)

Allen Wang's avatar
Allen Wang committed
417
  if encoder_type == "xlnet":
Frederick Liu's avatar
Frederick Liu committed
418
    return networks.XLNetBase(
Allen Wang's avatar
Allen Wang committed
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
        vocab_size=encoder_cfg.vocab_size,
        num_layers=encoder_cfg.num_layers,
        hidden_size=encoder_cfg.hidden_size,
        num_attention_heads=encoder_cfg.num_attention_heads,
        head_size=encoder_cfg.head_size,
        inner_size=encoder_cfg.inner_size,
        dropout_rate=encoder_cfg.dropout_rate,
        attention_dropout_rate=encoder_cfg.attention_dropout_rate,
        attention_type=encoder_cfg.attention_type,
        bi_data=encoder_cfg.bi_data,
        two_stream=encoder_cfg.two_stream,
        tie_attention_biases=encoder_cfg.tie_attention_biases,
        memory_length=encoder_cfg.memory_length,
        clamp_length=encoder_cfg.clamp_length,
        reuse_length=encoder_cfg.reuse_length,
        inner_activation=encoder_cfg.inner_activation,
        use_cls_mask=encoder_cfg.use_cls_mask,
        embedding_width=encoder_cfg.embedding_width,
        initializer=tf.keras.initializers.RandomNormal(
            stddev=encoder_cfg.initializer_range))

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
  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)

Hongkun Yu's avatar
Hongkun Yu committed
474
475
  # Uses the default BERTEncoder configuration schema to create the encoder.
  # If it does not match, please add a switch branch by the encoder type.
Frederick Liu's avatar
Frederick Liu committed
476
  return networks.BertEncoder(
Hongkun Yu's avatar
Hongkun Yu committed
477
478
479
480
481
482
483
484
485
486
      vocab_size=encoder_cfg.vocab_size,
      hidden_size=encoder_cfg.hidden_size,
      num_layers=encoder_cfg.num_layers,
      num_attention_heads=encoder_cfg.num_attention_heads,
      intermediate_size=encoder_cfg.intermediate_size,
      activation=tf_utils.get_activation(encoder_cfg.hidden_activation),
      dropout_rate=encoder_cfg.dropout_rate,
      attention_dropout_rate=encoder_cfg.attention_dropout_rate,
      max_sequence_length=encoder_cfg.max_position_embeddings,
      type_vocab_size=encoder_cfg.type_vocab_size,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
487
      initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
488
          stddev=encoder_cfg.initializer_range),
Frederick Liu's avatar
Frederick Liu committed
489
      output_range=encoder_cfg.output_range,
Hongkun Yu's avatar
Hongkun Yu committed
490
      embedding_width=encoder_cfg.embedding_size,
Chen Chen's avatar
Chen Chen committed
491
      embedding_layer=embedding_layer,
492
      return_all_encoder_outputs=encoder_cfg.return_all_encoder_outputs,
493
494
      dict_outputs=True,
      norm_first=encoder_cfg.norm_first)