def test_inference_model(self):
model = TFRemBertModel.from_pretrained("google/rembert")
input_ids = tf.constant([[312, 56498, 313, 2125, 313]])
segment_ids = tf.constant([[0, 0, 0, 1, 1]])
output = model(input_ids,
token_type_ids=segment_ids,
output_hidden_states=True)
hidden_size = 1152
expected_shape = [1, 5, hidden_size]
self.assertEqual(output["last_hidden_state"].shape, expected_shape)
expected_implementation = tf.constant([[
[0.0754, -0.2022, 0.1904],
[-0.3354, -0.3692, -0.4791],
[-0.2314, -0.6729, -0.0749],
[-0.0396, -0.3105, -0.4234],
[-0.1571, -0.0525, 0.5353],
]])
tf.debugging.assert_near(output["last_hidden_state"][:, :, :3],
expected_implementation,
atol=1e-4)
def test_model_from_pretrained(self):
model = TFRemBertModel.from_pretrained("google/rembert")
self.assertIsNotNone(model)
