def _truncate_segments(segments, max_num_tokens, rng):
  """
  Truncate sequence pair according to original BERT implementation:
  https://github.com/google-research/bert/blob/master/create_pretraining_data.py#L391
  """
  while True:
    if sum(len(s) for s in segments)<=max_num_tokens:
      break

    segments = sorted(segments, key=lambda s:len(s), reverse=True)
    trunc_tokens = segments[0]

    assert len(trunc_tokens) >= 1

    if rng.random() < 0.5:
      trunc_tokens.pop(0)
    else:
      trunc_tokens.pop()
  return segments