def build_vocab(sents_data,
                task_name,
                vocab_mode="downstream_vocab",
                vocab_path=None):
  """find all words in input corpus to build a vocabulary."""
  if vocab_mode == "bert_vocab":
    # Use a custom vocab to carry out filtering (such as BERT's word piece)
    with gfile.Open(vocab_path, "r") as f:
      vocab = f.read().strip().split("\n")
    # Filter out special tokens
    vocab = [x for x in vocab if x[0] != "[" and x[-1] != "]"]
    probs = [1.0 / len(vocab) for _ in vocab]

  elif vocab_mode == "full_corpus":
    # Use all words in a corpus of text to find out the vocabulary
    vocab = collections.Counter("\n".join(sents_data).split())
    vocab = [(k, v) for k, v in vocab.items()]
    vocab.sort(key=lambda x: x[1], reverse=True)
    vocab_total = sum([x[1] for x in vocab])

    probs = [float(x[1]) / vocab_total for x in vocab]
    vocab = [x[0] for x in vocab]

  elif "full_corpus_top_" in vocab_mode:
    full_vocab = collections.defaultdict(int)
    for sent in sents_data:
      for word in sent.split():
        full_vocab[word] += 1
    # Sort the vocabulary words according to their frequency
    full_vocab = sorted([(k, v) for k, v in full_vocab.items()],
                        key=lambda x: x[1],
                        reverse=True)
    # Take the top-k values from the vocabulary for the final list
    top_k_val = int(vocab_mode[len("full_corpus_top_"):])
    vocab = [x[0] for x in full_vocab[:top_k_val]]
    probs = [1.0 / len(vocab) for _ in vocab]

  elif vocab_mode == "downstream_vocab":
    vocab = collections.defaultdict(int)
    for sent in sents_data:
      for index in task_input_indices[task_name]:
        original_sent = sent.split("\t")[index].split()
        for word in original_sent:
          vocab[word] += 1

    vocab = [(k, v) for k, v in vocab.items()]
    vocab.sort(key=lambda x: x[1], reverse=True)
    vocab_total = sum([x[1] for x in vocab])

    probs = [float(x[1]) / vocab_total for x in vocab]
    vocab = [x[0] for x in vocab]

  else:
    probs = None
    vocab = None

  return vocab, probs
find all words in input corpus to build a vocabulary.
language/bert_extraction/steal_bert_classifier/data_generation/preprocess_util.py