def create_bert2bert_model(
params: configs.BERT2BERTConfig,
cls=Bert2Bert,
init_checkpoint: Optional[Text] = None) -> tf.keras.Model:
"""A helper to create Bert2Bert model."""
bert_layer, decoder_layer = get_bert2bert_layers(params=params)
if init_checkpoint:
utils.initialize_bert2bert_from_pretrained_bert(
bert_layer, decoder_layer, init_checkpoint)
return cls(params=params,
bert_layer=bert_layer,
decoder_layer=decoder_layer,
name="bert2bert")
def create_nhnet_model(
params: configs.NHNetConfig,
cls=NHNet,
init_checkpoint: Optional[Text] = None) -> tf.keras.Model:
"""A helper to create NHNet model."""
bert_layer, decoder_layer = get_nhnet_layers(params=params)
model = cls(params=params,
bert_layer=bert_layer,
decoder_layer=decoder_layer,
name="nhnet")
if init_checkpoint:
logging.info(
"Checkpoint file %s found and restoring from "
"initial checkpoint.", init_checkpoint)
if params.init_from_bert2bert:
ckpt = tf.train.Checkpoint(model=model)
ckpt.restore(init_checkpoint).assert_existing_objects_matched()
else:
utils.initialize_bert2bert_from_pretrained_bert(
bert_layer, decoder_layer, init_checkpoint)
return model