def truncate_segment_tokens(segment_tokens, max_tokens):
  """Truncates segment tokens.

  Suppose we arrange our segment tokens in a matrix, e.g.:

  this  is    query  one    <pad>       <-- first segment
  this  is    query  two    <pad>       <-- second segment
  more  text  <pad>  <pad>  <pad>       <-- third segment

  If we visit the tokens in column-major order, we get this ordering:

  1     4     7     9
  2     5     8     10
  3     6

  This function removes all tokens with visit index > max_tokens.

  This function generalizes language.bert.run_classifier._truncate_seq_pair,
  and exactly matches its behavior when num_segments = 2 (a pair).

  Args:
    segment_tokens (RaggedTensor): a 2-D RaggedTensor of strings. One row for
      each segment. Each row is a list of tokens.
    max_tokens (Tensor): scalar int, max number of tokens for all segments
      combined.

  Returns:
    segment_tokens_truncated (RaggedTensor)
  """
  # mask is a 2-D int32 Tensor: 1's indicate tokens, 0's indicate padding.
  mask = tf.ones_like(segment_tokens, dtype=tf.int32).to_tensor()
  mask.set_shape([None, None])

  max_tokens = tf.cast(max_tokens, dtype=tf.int32)

  # visit_order reflects the column-major order we would visit each entry.
  transposed_mask = tf.transpose(mask)
  visit_order_flat = tf.cumsum(tf.reshape(transposed_mask, [-1]))
  visit_order = tf.transpose(
      tf.reshape(visit_order_flat, tf.shape(transposed_mask)))

  should_keep = (visit_order <= max_tokens) & tf.cast(mask, tf.bool)
  truncated_row_lengths = tf.reduce_sum(tf.cast(should_keep, tf.int64), axis=1)
  tokens_to_keep = tf.boolean_mask(segment_tokens.to_tensor(), should_keep)
  return tf.RaggedTensor.from_row_lengths(tokens_to_keep, truncated_row_lengths)
Truncates segment tokens.

  Suppose we arrange our segment tokens in a matrix, e.g.:

  this  is    query  one    <pad>       <-- first segment
  this  is    query  two    <pad>       <-- second segment
  more  text  <pad>  <pad>  <pad>       <-- third segment

  If we visit the tokens in column-major order, we get this ordering:

  1     4     7     9
  2     5     8     10
  3     6

  This function removes all tokens with visit index > max_tokens.

  This function generalizes language.bert.run_classifier._truncate_seq_pair,
  and exactly matches its behavior when num_segments = 2 (a pair).

  Args:
    segment_tokens (RaggedTensor): a 2-D RaggedTensor of strings. One row for
      each segment. Each row is a list of tokens.
    max_tokens (Tensor): scalar int, max number of tokens for all segments
      combined.

  Returns:
    segment_tokens_truncated (RaggedTensor)
  
language/conpono/reconstruct/preprocess.py