def sample_next_sequence(vocab, probs, seq_length, scheme):
  """Sample a full sequence of random tokens using uniform or multinomials."""
  if scheme.endswith("_uniform"):
    return [random.choice(vocab) for _ in range(seq_length)]

  elif scheme.endswith("_freq"):
    bow_sent = np.random.multinomial(seq_length, probs)
    nonzero_indices = np.nonzero(bow_sent)[0]
    bow_indices = []
    for index, freq in zip(nonzero_indices, bow_sent[nonzero_indices]):
      bow_indices.extend([index for _ in range(freq)])
    assert len(bow_indices) == seq_length
    random.shuffle(bow_indices)
    return [vocab[i] for i in bow_indices]

  elif scheme.endswith("_freq_thresh"):
    freq_seq_length = int(np.sum(probs[:10000]) * seq_length)
    # use frequency sampling for first freq_seq_length tokens
    bow_indices = []
    if freq_seq_length > 0:
      bow_sent = np.random.multinomial(seq_length, probs[:10000])
      nonzero_indices = np.nonzero(bow_sent)[0]
      for index, freq in zip(nonzero_indices, bow_sent[nonzero_indices]):
        bow_indices.extend([index for _ in range(freq)])

    # Fill in remaining tokens with uniform random indices
    while len(bow_indices) < seq_length:
      bow_indices.append(random.randint(0, len(vocab) - 1))

    random.shuffle(bow_indices)
    return [vocab[i] for i in bow_indices]
  else:
    return None
Sample a full sequence of random tokens using uniform or multinomials.
language/bert_extraction/steal_bert_classifier/data_generation/preprocess_util.py