def build_bert_inputs(example):
  """Convert example <Tensor [30, 70]> into bert inputs."""
  k_size = FLAGS.k_size

  CLS_ID = tf.constant([101], dtype=tf.int64)  # pylint: disable=invalid-name
  SEP_ID = tf.constant([102], dtype=tf.int64)  # pylint: disable=invalid-name
  max_len = tf.constant([FLAGS.max_para_length])
  context_size = tf.constant([FLAGS.context_size])

  intermediate_examples_tensor = tf.reduce_sum(tf.abs(example), 1)
  examples_zero_vector = tf.zeros(shape=(1, 1), dtype=tf.int64)
  examples_bool_mask = tf.squeeze(
      tf.not_equal(intermediate_examples_tensor, examples_zero_vector))
  paragraph_len = tf.reduce_sum(tf.cast(examples_bool_mask, tf.int32))

  start = tf.random.uniform([1],
                            0,
                            tf.reshape(paragraph_len, []) -
                            tf.reshape(context_size, []) + 1,
                            dtype=tf.int32)

  # Slice the document into the before, after and context.
  # Discard the zero padding.
  sizes = tf.squeeze(
      tf.concat([[
          start, context_size, paragraph_len - context_size - start,
          max_len - paragraph_len
      ]], 0))
  before, context, after, _ = tf.split(example, sizes, axis=0)

  # Gather the context removing zero padding at end of sentences.
  non_zeros = tf.where(tf.not_equal(context, tf.zeros_like(context)))
  context_gathered = tf.gather_nd(context, non_zeros)

  # Flip before so we select the 4 sentences closest to target
  before = tf.reverse(before, axis=[0])

  # pad both to longer than needed
  paddings = tf.constant([[0, 8], [0, 0]])
  before = tf.pad(before, paddings)
  after = tf.pad(after, paddings)

  # Extend targets to 3 sentences
  # pad both
  before_minus_one = before[1:][:k_size]
  before_minus_two = before[2:][:k_size]
  after_plus_one = after[1:][:k_size]
  after_plus_two = after[2:][:k_size]
  before = before[:k_size]
  after = after[:k_size]

  before = tf.concat([before_minus_two, before_minus_one, before], axis=1)
  after = tf.concat([after, after_plus_one, after_plus_two], axis=1)
  ############################################################################

  # These 8 sentences are the 8 surrounding targets. Some are padding.
  targets = tf.concat([before, after], axis=0)

  # Remove the padding from the sourrounding sentences
  # Eg. if context starts at beginning of paragraph, before is all padding
  intermediate_tensor = tf.reduce_sum(tf.abs(targets), 1)
  zero_vector = tf.zeros(shape=(1, 1), dtype=tf.int64)
  bool_mask = tf.squeeze(tf.not_equal(intermediate_tensor, zero_vector))
  bool_mask.set_shape([None])
  targets = tf.boolean_mask(targets, bool_mask)

  # Randomly select 4 targets
  # We will also select the label_types for each selected target
  indices = tf.range(0, limit=tf.shape(targets)[0], dtype=tf.int32)
  shuffled_indices = tf.random.shuffle(indices)[:k_size]

  targets = tf.gather(targets, shuffled_indices)
  if k_size == 4:
    full_labels = tf.concat([tf.range(3, -1, -1), tf.range(4, 8)], axis=0)
  elif k_size == 3:
    full_labels = tf.concat([tf.range(2, -1, -1), tf.range(3, 6)], axis=0)
  elif k_size == 2:
    full_labels = tf.concat([tf.range(1, -1, -1), tf.range(2, 4)], axis=0)
  elif k_size == 1:
    full_labels = tf.concat([tf.range(0, -1, -1), tf.range(1, 2)], axis=0)
  label_types = tf.boolean_mask(full_labels, bool_mask)
  label_types = tf.gather(label_types, shuffled_indices)

  # create inputs
  bert_inputs = []
  input_masks = []
  segment_ids = []

  # make context
  ctx_segment_id = tf.concat([
      tf.zeros_like(CLS_ID, dtype=tf.int64),
      tf.zeros_like(context_gathered),
      tf.zeros_like(SEP_ID, dtype=tf.int64)
  ],
                             axis=0)
  ctx_segment_id = pad_and_cut(ctx_segment_id, FLAGS.max_seq_length)
  segment_ids.append(ctx_segment_id)

  new_ctx_input = tf.concat([CLS_ID, context_gathered, SEP_ID], axis=0)
  ctx_input_mask = tf.ones_like(new_ctx_input)
  ctx_input_mask = pad_and_cut(ctx_input_mask, FLAGS.max_seq_length)
  input_masks.append(ctx_input_mask)
  padded_new_ctx_input = pad_and_cut(new_ctx_input, FLAGS.max_seq_length)
  bert_inputs.append(padded_new_ctx_input)

  for i in range(k_size):
    target_non_zero = tf.where(
        tf.not_equal(targets[i], tf.zeros_like(targets[i])))
    targets_stripped = tf.gather_nd(targets[i], target_non_zero)
    if FLAGS.include_context:
      segment_id = tf.concat([
          tf.zeros_like(CLS_ID, dtype=tf.int64),
          tf.zeros_like(context_gathered),
          tf.zeros_like(SEP_ID, dtype=tf.int64),
          tf.ones_like(targets_stripped),
          tf.ones_like(SEP_ID, dtype=tf.int64)
      ],
                             axis=0)
    else:
      segment_id = tf.concat([
          tf.zeros_like(CLS_ID, dtype=tf.int64),
          tf.zeros_like(targets_stripped),
          tf.zeros_like(SEP_ID, dtype=tf.int64)
      ],
                             axis=0)
    segment_id = pad_and_cut(segment_id, FLAGS.max_seq_length)
    segment_ids.append(segment_id)
    if FLAGS.include_context:
      new_input = tf.concat(
          [CLS_ID, context_gathered, SEP_ID, targets_stripped, SEP_ID], axis=0)
    else:
      new_input = tf.concat([CLS_ID, targets_stripped, SEP_ID], axis=0)
    input_mask = tf.ones_like(new_input)
    input_mask = pad_and_cut(input_mask, FLAGS.max_seq_length)
    input_masks.append(input_mask)
    padded_new_input = pad_and_cut(new_input, FLAGS.max_seq_length)
    bert_inputs.append(padded_new_input)
  bert_inputs = tf.stack(bert_inputs, axis=0)
  input_masks = tf.stack(input_masks, axis=0)
  segment_ids = tf.stack(segment_ids, axis=0)

  out = Outputs_And_Context(bert_inputs, input_masks, segment_ids, label_types,
                            context_gathered)

  return out
Convert example <Tensor [30, 70]> into bert inputs.
language/conpono/cpc/run_bilin_cpc.py