def _truncate_and_pad_token_ids(token_ids, max_length):
  """Truncates or pads the token id list to max length."""
  token_ids = token_ids[:max_length]
  padding_size = max_length - len(token_ids)
  if padding_size > 0:
    token_ids += [0] * padding_size
  return token_ids