def test_create_embedding_table_for_slots(self): slot_name = "momentum" init_value = 3.5 table = EmbeddingTable( get_slot_table_name(self.name, slot_name), dim=self.dim, initializer=init_value, is_slot=True, ) self.assertIsNotNone(table) self.assertEqual(table.name, get_slot_table_name(self.name, slot_name)) self.assertEqual(table.dim, self.dim) # test initialize embedding = table.get([2]) self.assertTrue((embedding - init_value < 0.0001).all())
class EmbeddingTableTest(unittest.TestCase): def setUp(self): self.name = "embedding_1" self.dim = 10 self.initializer = "uniform" self.table = EmbeddingTable(self.name, self.dim, self.initializer) def test_embedding_table_init(self): self.assertIsNotNone(self.table) self.assertEqual(self.table.name, self.name) self.assertEqual(self.table.dim, self.dim) self.assertEqual( tf.keras.initializers.get(self.initializer).__class__, self.table.initializer.__class__, ) def test_embedding_table_get(self): self.table.clear() indices = [0, 3, 7] res = self.table.get(indices) self.assertTupleEqual(res.shape, (3, 10)) res = self.table.get([]) self.assertIsNone(res) self.table.get([0, 3, 8]) self.assertEqual(len(self.table.embedding_vectors), 4) def test_embedding_table_set(self): self.table.clear() indices = [0, 1, 4] x = len(indices) values = np.random.uniform(size=x * self.dim).reshape((x, self.dim)) self.table.set(indices, values) row0 = self.table.get([0]) row1 = self.table.get([1]) row4 = self.table.get([4]) rows = [row0, row1, row4] rows = np.concatenate(rows) np.testing.assert_array_equal(rows, values) def test_create_embedding_table(self): embedding_pb = EmbeddingTableInfo() embedding_pb.name = self.name embedding_pb.dim = self.dim embedding_pb.initializer = self.initializer table = create_embedding_table(embedding_pb) self.assertIsNotNone(table) self.assertEqual(table.name, self.name) self.assertEqual( tf.keras.initializers.get(self.initializer).__class__, table.initializer.__class__, ) self.assertEqual(table.dim, self.dim) def test_create_embedding_table_for_slots(self): slot_name = "momentum" init_value = 3.5 table = EmbeddingTable( get_slot_table_name(self.name, slot_name), dim=self.dim, initializer=init_value, is_slot=True, ) self.assertIsNotNone(table) self.assertEqual(table.name, get_slot_table_name(self.name, slot_name)) self.assertEqual(table.dim, self.dim) # test initialize embedding = table.get([2]) self.assertTrue((embedding - init_value < 0.0001).all())