import tensorflow as tf


    vocab_path = self.vocabulary_file_by_name(vocab_filename)
    if not vocab_path:
      raise ValueError(
          'Could not compute vocabulary size for {}, does not exist'.format(
              vocab_filename))
    elif vocab_path.endswith('tfrecord.gz'):
      dataset = tf.data.TFRecordDataset(vocab_path, compression_type='GZIP')

      def reduce_fn(accum, elem):
        return tf.size(elem, out_type=tf.int64, name='vocabulary_size') + accum

      return _get_tensor_value(
          dataset.batch(tf.int32.max).reduce(
              tf.constant(0, tf.int64), reduce_fn))
    else:
      with tf.io.gfile.GFile(vocab_path, 'rb') as f:
        return sum(1 for _ in f)

  def vocabulary_by_name(self, vocab_filename: str) -> List[bytes]:
