Ejemplo n.º 1
0
 def test_no_segment_ids(self):
   segment_matrix = xlnet_base._compute_segment_matrix(
       segment_ids=None,
       memory_length=2,
       batch_size=1,
       use_cls_mask=False)
   self.assertIsNone(segment_matrix)
Ejemplo n.º 2
0
 def test_basic(self):
     batch_size = 1
     memory_length = 0
     segment_ids = np.array([[1, 1, 2, 1]])
     expected_segment_matrix = np.array([[[False, False, True, False],
                                          [False, False, True, False],
                                          [True, True, False, True],
                                          [False, False, True, False]]])
     segment_matrix = xlnet_base._compute_segment_matrix(
         segment_ids=segment_ids,
         memory_length=memory_length,
         batch_size=batch_size,
         use_cls_mask=False)
     self.assertAllClose(segment_matrix, expected_segment_matrix)
Ejemplo n.º 3
0
 def test_basic_with_memory(self):
     batch_size = 1
     memory_length = 1
     segment_ids = np.array([[1, 1, 2, 1]])
     expected_segment_matrix = np.array([[[True, False, False, True, False],
                                          [True, False, False, True, False],
                                          [True, True, True, False, True],
                                          [True, False, False, True,
                                           False]]]).astype(int)
     segment_matrix = tf.cast(xlnet_base._compute_segment_matrix(
         segment_ids=segment_ids,
         memory_length=memory_length,
         batch_size=batch_size,
         use_cls_mask=False),
                              dtype=tf.uint8)
     self.assertAllClose(segment_matrix, expected_segment_matrix)
Ejemplo n.º 4
0
 def dont_test_basic_with_class_mask(self):
     # TODO(allencwang) - this test should pass but illustrates the legacy issue
     # of using class mask. Enable once addressed.
     batch_size = 1
     memory_length = 0
     segment_ids = np.array([[1, 1, 2, 1]])
     expected_segment_matrix = np.array([[[False, False, True, False],
                                          [False, False, True, False],
                                          [True, True, False, True],
                                          [False, False, True,
                                           False]]]).astype(int)
     segment_matrix = tf.cast(xlnet_base._compute_segment_matrix(
         segment_ids=segment_ids,
         memory_length=memory_length,
         batch_size=batch_size,
         use_cls_mask=True),
                              dtype=tf.uint8)
     self.assertAllClose(segment_matrix, expected_segment_matrix)