def test_export_saved_model(self):
model = Bert(
vocab_size=21128, num_layers=12, num_attention_heads=8, return_states=True, return_attention_weights=True
)
input_ids, segment_ids, input_mask = model.dummy_inputs()
model(inputs=[input_ids, segment_ids, input_mask])
model.summary()
model.save("models/export/2", include_optimizer=False)
def test_build_model(self):
model = Bert(vocab_size=21128)
input_ids, segment_ids, input_mask = model.dummy_inputs()
model(inputs=[input_ids, segment_ids, input_mask])
model.summary()