def test_cut_or_padding(self):
# test for 1d
origin_1_t = tf.placeholder(dtype=tf.int32, shape=[None])
after_1_t = cut_or_padding(origin_1_t, 3)
with self.cached_session(use_gpu=False, force_gpu=False) as sess:
# test for padding
res = sess.run(after_1_t, feed_dict={origin_1_t: [1, 2]})
self.assertAllEqual(res, [1, 2, 0])
# test for cut
res = sess.run(after_1_t, feed_dict={origin_1_t: [1, 2, 3, 4, 5]})
self.assertAllEqual(res, [1, 2, 3])
# test for 2d
origin_2_t = tf.placeholder(dtype=tf.int32, shape=[None, None])
after_2_t = cut_or_padding(origin_2_t, 3)
with self.cached_session(use_gpu=False, force_gpu=False) as sess:
# test for padding
res = sess.run(after_2_t, feed_dict={origin_2_t: [[1, 2], [1, 2]]})
self.assertAllEqual(res, [[1, 2, 0], [1, 2, 0]])
# test for cut
res = sess.run(
after_2_t,
feed_dict={origin_2_t: [[1, 2, 3, 4], [1, 2, 3, 4]]})
self.assertAllEqual(res, [[1, 2, 3], [1, 2, 3]])
def pad_to_hier_input(inputs, max_doc_len, max_sen_len, padding_token=0): """ Input shape: [batch_size, max_len] New Input shape: [batch_size, max_doc_len, max_sen_len] """ new_len = max_sen_len * max_doc_len new_input = cut_or_padding(inputs, new_len, padding_token=padding_token) new_input = tf.reshape(new_input, [-1, max_doc_len, max_sen_len]) return new_input