import tensorflow as tf

    elif vocab_path.endswith('tfrecord.gz'):
      dataset = tf.data.TFRecordDataset(vocab_path, compression_type='GZIP')
      vocab_tensor = dataset.batch(tf.int32.max).reduce(
          tf.constant([], dtype=tf.string),
          lambda state, elem: tf.concat([state, elem], axis=-1))
      # Using as_numpy_iterator only works when executing eagerly.
