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 _check_bert_outputs(self, return_states=False, return_attention_weights=False):
NUM_LAYERS = 4
model = Bert(
vocab_size=100,
num_layers=NUM_LAYERS,
return_states=return_states,
return_attention_weights=return_attention_weights,
)
input_ids, segment_ids, attn_mask = self._build_bert_inputs()
outputs = model(inputs=[input_ids, segment_ids, attn_mask])
sequence_outputs, pooled_outputs = outputs[0], outputs[1]
self.assertAllEqual([2, 16, 768], sequence_outputs.shape)
self.assertAllEqual([2, 768], pooled_outputs.shape)
all_states, all_attn_weights = None, None
if return_states and return_attention_weights:
self.assertEqual(4, len(outputs))
all_states, all_attn_weights = outputs[2], outputs[3]
elif return_states and not return_attention_weights:
self.assertEqual(3, len(outputs))
all_states = outputs[2]
elif not return_states and return_attention_weights:
self.assertEqual(3, len(outputs))
all_attn_weights = outputs[2]
else:
self.assertEqual(2, len(outputs))
if all_states is not None:
# self.assertEqual(2, len(all_states))
# for state in all_states:
# self.assertAllEqual([2, 16, 768], state.shape)
self.assertAllEqual([2, NUM_LAYERS, 16, 768], all_states.shape)
if all_attn_weights is not None:
# self.assertEqual(2, len(all_attn_weights))
# for attention in all_attn_weights:
# self.assertAllEqual([2, 8, 16, 16], attention.shape)
self.assertAllEqual([2, NUM_LAYERS, 8, 16, 16], all_attn_weights.shape)
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()
def test_bert_config(self):
model = Bert(vocab_size=100, num_layers=2, return_states=True, return_attention_weights=True)
config = model.get_config()
print(config)