def _get_rc_model_input(
    question_ids,
    question_mask,
    context_ids,
    context_mask,
    vocab,
):
  """Create RC module input from separate batched components.

  Args:
    question_ids: <int32> [batch_size, question_len]
    question_mask: <int32> [batch_size, question_len]
    context_ids: <int32> [batch_size, context_len]
    context_mask: <int32> [batch_size, context_len]
    vocab: Instance of text_utils.Vocab.

  Returns:
    input_ids: <int32> [batch_size, rc_input_len]
    input_mask: <int32> [batch_size, rc_input_len]
    segment_ids: <int32> [batch_size, rc_input_len]
  """
  # Get batch size.
  batch_size = tensor_utils.shape(context_ids, 0)

  # Get special tokens.
  cls = vocab.t2i(vocab.CLS)
  sep = vocab.t2i(vocab.SEP)

  # Join question, context, and special tokens.
  cls_batch = tf.fill([batch_size, 1], cls)
  sep_batch = tf.fill([batch_size, 1], sep)
  input_ids = tf.concat(
      [cls_batch, question_ids, sep_batch, context_ids, sep_batch], axis=1)

  # Create and join segment ids.
  segment_a_ids = tf.fill(
      [batch_size, tensor_utils.shape(question_ids, 1) + 2], 0)
  segment_b_ids = tf.fill(
      [batch_size, tensor_utils.shape(context_ids, 1) + 1], 1)
  segment_ids = tf.concat([segment_a_ids, segment_b_ids], axis=1)

  # Create joined mask, accounting for special tokens gaps.
  gap_mask = tf.fill([batch_size, 1], 1)
  input_mask = tf.concat(
      [gap_mask, question_mask, gap_mask, context_mask, gap_mask], axis=1)
  bool_mask = tf.cast(input_mask, tf.bool)

  # Select unmasked items and move all padding to the end.
  # Right now this looks like this:
  #   [CLS] X X X [PAD] ... [SEP] Y Y Y [PAD] ... [SEP] [PAD] ...
  # And we want to change it to look like this:
  #   [CLS] X X X [SEP] Y Y Y [SEP] [PAD] ...
  input_ids = tensor_utils.boolean_mask(input_ids, bool_mask)
  input_mask = tensor_utils.boolean_mask(input_mask, bool_mask)
  segment_ids = tensor_utils.boolean_mask(segment_ids, bool_mask)

  return input_ids, input_mask, segment_ids
Create RC module input from separate batched components.

  Args:
    question_ids: <int32> [batch_size, question_len]
    question_mask: <int32> [batch_size, question_len]
    context_ids: <int32> [batch_size, context_len]
    context_mask: <int32> [batch_size, context_len]
    vocab: Instance of text_utils.Vocab.

  Returns:
    input_ids: <int32> [batch_size, rc_input_len]
    input_mask: <int32> [batch_size, rc_input_len]
    segment_ids: <int32> [batch_size, rc_input_len]
  
language/capwap/utils/reward_utils.py