Skip to content

GitLab

Commit dd0126f9 authored by Chen Chen's avatar Chen Chen Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 328798031
parent 20e2cb97
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 2 deletions
official/nlp/configs/encoders.py
View file @ dd0126f9
...@@ -45,6 +45,7 @@ class BertEncoderConfig(hyperparams.Config): ...@@ -45,6 +45,7 @@ class BertEncoderConfig(hyperparams.Config):
type_vocab_size: int = 2 type_vocab_size: int = 2
initializer_range: float = 0.02 initializer_range: float = 0.02
embedding_size: Optional[int] = None embedding_size: Optional[int] = None
return_all_encoder_outputs: bool = False
@dataclasses.dataclass @dataclasses.dataclass
...@@ -186,7 +187,8 @@ def build_encoder(config: EncoderConfig, ...@@ -186,7 +187,8 @@ def build_encoder(config: EncoderConfig,
num_hidden_instances=encoder_cfg.num_layers, num_hidden_instances=encoder_cfg.num_layers,
pooled_output_dim=encoder_cfg.hidden_size, pooled_output_dim=encoder_cfg.hidden_size,
pooler_layer_initializer=tf.keras.initializers.TruncatedNormal( pooler_layer_initializer=tf.keras.initializers.TruncatedNormal(
stddev=encoder_cfg.initializer_range)) stddev=encoder_cfg.initializer_range),
return_all_layer_outputs=encoder_cfg.return_all_encoder_outputs)
return encoder_cls(**kwargs) return encoder_cls(**kwargs)
if encoder_type == "mobilebert": if encoder_type == "mobilebert":
...@@ -242,4 +244,5 @@ def build_encoder(config: EncoderConfig, ...@@ -242,4 +244,5 @@ def build_encoder(config: EncoderConfig,
initializer=tf.keras.initializers.TruncatedNormal( initializer=tf.keras.initializers.TruncatedNormal(
stddev=encoder_cfg.initializer_range), stddev=encoder_cfg.initializer_range),
embedding_width=encoder_cfg.embedding_size, embedding_width=encoder_cfg.embedding_size,
embedding_layer=embedding_layer) embedding_layer=embedding_layer,
return_all_encoder_outputs=encoder_cfg.return_all_encoder_outputs)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment