def test_test_preprocess_fn_return_dataset_element_spec(self): token = collections.OrderedDict(tokens=([ 'one must imagine', ])) ds = tf.data.Dataset.from_tensor_slices(token) test_preprocess_fn = dataset.create_test_dataset_preprocess_fn( max_seq_len=10, vocab=['one', 'must']) test_preprocessed_ds = test_preprocess_fn(ds) self.assertEqual(test_preprocessed_ds.element_spec, (tf.TensorSpec(shape=[None, 10], dtype=tf.int64), tf.TensorSpec(shape=[None, 10], dtype=tf.int64)))
def test_test_preprocess_fn_returns_correct_sequence(self): token = collections.OrderedDict(tokens=([ 'one must imagine', ])) ds = tf.data.Dataset.from_tensor_slices(token) test_preprocess_fn = dataset.create_test_dataset_preprocess_fn( max_seq_len=6, vocab=['one', 'must']) test_preprocessed_ds = test_preprocess_fn(ds) element = next(iter(test_preprocessed_ds)) # BOS is len(vocab)+2, EOS is len(vocab)+3, pad is 0, OOV is len(vocab)+1 self.assertAllEqual(self.evaluate(element[0]), np.array([[4, 1, 2, 3, 5, 0]]))