Generate subtoken vocabulary close to the target size.

def _generate_subtokens_with_target_vocab_size(
    token_counts, alphabet, target_size, threshold, min_count=None,
    reserved_tokens=None):
  """Generate subtoken vocabulary close to the target size."""
  if reserved_tokens is None:
    reserved_tokens = RESERVED_TOKENS

  if min_count is not None:
    tf.logging.info("Using min_count=%d to generate vocab with target size %d" %
                    (min_count, target_size))
    return _generate_subtokens(
        token_counts, alphabet, min_count, reserved_tokens=reserved_tokens)

  def bisect(min_val, max_val):
    """Recursive function to binary search for subtoken vocabulary."""
    cur_count = (min_val + max_val) // 2
    tf.logging.info("Binary search: trying min_count=%d (%d %d)" %
                    (cur_count, min_val, max_val))
    subtoken_list = _generate_subtokens(
        token_counts, alphabet, cur_count, reserved_tokens=reserved_tokens)

    val = len(subtoken_list)
    tf.logging.info("Binary search: min_count=%d resulted in %d tokens" %
                    (cur_count, val))

    within_threshold = abs(val - target_size) < threshold
    if within_threshold or min_val >= max_val or cur_count < 2:
      return subtoken_list
    if val > target_size:
      other_subtoken_list = bisect(cur_count + 1, max_val)
    else:
      other_subtoken_list = bisect(min_val, cur_count - 1)

    # Return vocabulary dictionary with the closest number of tokens.
    other_val = len(other_subtoken_list)
    if abs(other_val - target_size) < abs(val - target_size):
      return other_subtoken_list
    return subtoken_list

  tf.logging.info("Finding best min_count to get target size of %d" %
                  target_size)
  return bisect(_MIN_MIN_COUNT, _MAX_MIN_COUNT)