encoders.py 21.6 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


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
198
@dataclasses.dataclass
class ReuseEncoderConfig(hyperparams.Config):
  """Reuse 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
  max_position_embeddings: int = 512
  type_vocab_size: int = 2
  initializer_range: float = 0.02
  embedding_size: Optional[int] = None
  output_range: Optional[int] = None
  return_all_encoder_outputs: bool = False
  # Pre/Post-LN Transformer
  norm_first: bool = False
  # Reuse transformer
  reuse_attention: int = -1
  use_relative_pe: bool = False
  pe_max_seq_length: int = 512
  max_reuse_layer_idx: int = 6


Allen Wang's avatar
Allen Wang committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
@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
224
225
226
227
@dataclasses.dataclass
class EncoderConfig(hyperparams.OneOfConfig):
  """Encoder configuration."""
  type: Optional[str] = "bert"
Chen Chen's avatar
Chen Chen committed
228
  albert: AlbertEncoderConfig = AlbertEncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
229
  bert: BertEncoderConfig = BertEncoderConfig()
Frederick Liu's avatar
Frederick Liu committed
230
  bert_v2: BertEncoderConfig = BertEncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
231
  bigbird: BigBirdEncoderConfig = BigBirdEncoderConfig()
Frederick Liu's avatar
Frederick Liu committed
232
  kernel: KernelEncoderConfig = KernelEncoderConfig()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
233
  mobilebert: MobileBertEncoderConfig = MobileBertEncoderConfig()
234
  reuse: ReuseEncoderConfig = ReuseEncoderConfig()
Allen Wang's avatar
Allen Wang committed
235
  xlnet: XLNetEncoderConfig = XLNetEncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
236
237
  # If `any` is used, the encoder building relies on any.BUILDER.
  any: hyperparams.Config = hyperparams.Config()
Hongkun Yu's avatar
Hongkun Yu committed
238
239
240


@gin.configurable
Hongkun Yu's avatar
Hongkun Yu committed
241
242
243
244
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
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
  """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
261
262
263
  encoder_type = config.type
  encoder_cfg = config.get()
  if encoder_cls and encoder_cls.__name__ == "EncoderScaffold":
Hongkun Yu's avatar
Hongkun Yu committed
264
    embedding_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
265
266
267
268
        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
269
        initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
270
271
            stddev=encoder_cfg.initializer_range),
        dropout_rate=encoder_cfg.dropout_rate,
Hongkun Yu's avatar
Hongkun Yu committed
272
273
    )
    hidden_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
274
275
        num_attention_heads=encoder_cfg.num_attention_heads,
        intermediate_size=encoder_cfg.intermediate_size,
Hongkun Yu's avatar
Hongkun Yu committed
276
        intermediate_activation=tf_utils.get_activation(
Hongkun Yu's avatar
Hongkun Yu committed
277
278
279
            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
280
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
281
            stddev=encoder_cfg.initializer_range),
Hongkun Yu's avatar
Hongkun Yu committed
282
283
284
285
    )
    kwargs = dict(
        embedding_cfg=embedding_cfg,
        hidden_cfg=hidden_cfg,
Hongkun Yu's avatar
Hongkun Yu committed
286
287
        num_hidden_instances=encoder_cfg.num_layers,
        pooled_output_dim=encoder_cfg.hidden_size,
Hongkun Yu's avatar
Hongkun Yu committed
288
        pooler_layer_initializer=tf.keras.initializers.TruncatedNormal(
Chen Chen's avatar
Chen Chen committed
289
            stddev=encoder_cfg.initializer_range),
290
291
        return_all_layer_outputs=encoder_cfg.return_all_encoder_outputs,
        dict_outputs=True)
Hongkun Yu's avatar
Hongkun Yu committed
292
293
    return encoder_cls(**kwargs)

Hongkun Yu's avatar
Hongkun Yu committed
294
295
296
297
298
299
300
301
302
303
  if encoder_type == "any":
    encoder = encoder_cfg.BUILDER(encoder_cfg)
    if not isinstance(encoder,
                      (tf.Module, tf.keras.Model, tf.keras.layers.Layer)):
      raise ValueError("The BUILDER returns an unexpected instance. The "
                       "`build_encoder` should returns a tf.Module, "
                       "tf.keras.Model or tf.keras.layers.Layer. However, "
                       f"we get {encoder.__class__}")
    return encoder

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
304
  if encoder_type == "mobilebert":
Frederick Liu's avatar
Frederick Liu committed
305
    return networks.MobileBERTEncoder(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
306
307
308
309
310
311
312
313
        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
314
        intermediate_act_fn=encoder_cfg.hidden_activation,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
315
316
317
        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
318
        initializer_range=encoder_cfg.initializer_range,
Chen Chen's avatar
Chen Chen committed
319
        use_bottleneck_attention=encoder_cfg.use_bottleneck_attention,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
320
321
322
        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
323
324
        classifier_activation=encoder_cfg.classifier_activation,
        input_mask_dtype=encoder_cfg.input_mask_dtype)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
325

Chen Chen's avatar
Chen Chen committed
326
  if encoder_type == "albert":
Frederick Liu's avatar
Frederick Liu committed
327
    return networks.AlbertEncoder(
Chen Chen's avatar
Chen Chen committed
328
329
330
331
332
333
334
335
336
337
338
339
        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(
340
341
            stddev=encoder_cfg.initializer_range),
        dict_outputs=True)
Chen Chen's avatar
Chen Chen committed
342

Hongkun Yu's avatar
Hongkun Yu committed
343
  if encoder_type == "bigbird":
Hongkun Yu's avatar
Hongkun Yu committed
344
345
    # TODO(frederickliu): Support use_gradient_checkpointing and update
    # experiments to use the EncoderScaffold only.
Frederick Liu's avatar
Frederick Liu committed
346
    if encoder_cfg.use_gradient_checkpointing:
Hongkun Yu's avatar
Hongkun Yu committed
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
      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
364
    embedding_cfg = dict(
Hongkun Yu's avatar
Hongkun Yu committed
365
        vocab_size=encoder_cfg.vocab_size,
Frederick Liu's avatar
Frederick Liu committed
366
        type_vocab_size=encoder_cfg.type_vocab_size,
Hongkun Yu's avatar
Hongkun Yu committed
367
        hidden_size=encoder_cfg.hidden_size,
Frederick Liu's avatar
Frederick Liu committed
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
        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
383
384
        num_attention_heads=encoder_cfg.num_attention_heads,
        intermediate_size=encoder_cfg.intermediate_size,
Frederick Liu's avatar
Frederick Liu committed
385
386
        intermediate_activation=tf_utils.get_activation(
            encoder_cfg.hidden_activation),
Hongkun Yu's avatar
Hongkun Yu committed
387
388
        dropout_rate=encoder_cfg.dropout_rate,
        attention_dropout_rate=encoder_cfg.attention_dropout_rate,
389
        norm_first=encoder_cfg.norm_first,
Frederick Liu's avatar
Frederick Liu committed
390
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
391
            stddev=encoder_cfg.initializer_range),
392
        attention_cls=layers.BigBirdAttention,
Frederick Liu's avatar
Frederick Liu committed
393
394
395
396
397
398
        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,
399
        mask_cls=layers.BigBirdMasks,
Frederick Liu's avatar
Frederick Liu committed
400
401
402
403
404
405
406
407
        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
408

Frederick Liu's avatar
Frederick Liu committed
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
  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
428
        scale=encoder_cfg.scale,
Frederick Liu's avatar
Frederick Liu committed
429
430
431
432
433
434
435
436
        )
    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,
437
        norm_first=encoder_cfg.norm_first,
Frederick Liu's avatar
Frederick Liu committed
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
        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
456
  if encoder_type == "xlnet":
Frederick Liu's avatar
Frederick Liu committed
457
    return networks.XLNetBase(
Allen Wang's avatar
Allen Wang committed
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
        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))

479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
  if encoder_type == "reuse":
    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)
    hidden_cfg = dict(
        num_attention_heads=encoder_cfg.num_attention_heads,
        inner_dim=encoder_cfg.intermediate_size,
        inner_activation=tf_utils.get_activation(
            encoder_cfg.hidden_activation),
        output_dropout=encoder_cfg.dropout_rate,
        attention_dropout=encoder_cfg.attention_dropout_rate,
        norm_first=encoder_cfg.norm_first,
        kernel_initializer=tf.keras.initializers.TruncatedNormal(
            stddev=encoder_cfg.initializer_range),
        reuse_attention=encoder_cfg.reuse_attention,
        use_relative_pe=encoder_cfg.use_relative_pe,
        pe_max_seq_length=encoder_cfg.pe_max_seq_length,
        max_reuse_layer_idx=encoder_cfg.max_reuse_layer_idx)
    kwargs = dict(
        embedding_cfg=embedding_cfg,
        hidden_cls=layers.ReuseTransformer,
        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=False,
        dict_outputs=True,
        feed_layer_idx=True,
        recursive=True)
    return networks.EncoderScaffold(**kwargs)

Frederick Liu's avatar
Frederick Liu committed
516
517
518
519
  bert_encoder_cls = networks.BertEncoder
  if encoder_type == "bert_v2":
    bert_encoder_cls = networks.BertEncoderV2

Hongkun Yu's avatar
Hongkun Yu committed
520
521
  # 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
522
  return bert_encoder_cls(
Hongkun Yu's avatar
Hongkun Yu committed
523
524
525
526
527
528
529
530
531
532
      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
533
      initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
534
          stddev=encoder_cfg.initializer_range),
Frederick Liu's avatar
Frederick Liu committed
535
      output_range=encoder_cfg.output_range,
Hongkun Yu's avatar
Hongkun Yu committed
536
      embedding_width=encoder_cfg.embedding_size,
Chen Chen's avatar
Chen Chen committed
537
      embedding_layer=embedding_layer,
538
      return_all_encoder_outputs=encoder_cfg.return_all_encoder_outputs,
539
540
      dict_outputs=True,
      norm_first=encoder_cfg.norm_first)