def output_order_with_beam(paragraphs, model_size, tokenizer):
  """Compute the predicted output for each paragraph."""

  tf.reset_default_graph()
  placeholders, model = create_cpc_model_and_placeholders(model_size)
  with tf.compat.v1.Session() as sess:
    _restore_checkpoint(FLAGS.model_weights)
    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())
    sess.run(init_op)

    # start with the first sentence
    # compute order greedily
    output_order_list = []
    for para_idx, paragraph in enumerate(paragraphs):
      beam_states = []
      tmp_softmax_mask = [True] * (len(paragraph) - 1) + [False] * (
          model_size + 1 - len(paragraph))
      tmp_output_order = []
      tmp_context = range(1)
      beam_states.append(
          State(1, tmp_output_order, tmp_softmax_mask, tmp_context))
      next_beam_states = []
      for _ in range(1, len(paragraph)):
        for prev_state in beam_states:
          cc = " ".join([paragraph[i] for i in prev_state.context])
          e = create_cpc_input_from_text(
              tokenizer, cc, paragraph[1:], [], group_size=model_size)
          while len(tokenizer.tokenize(cc)) > 400:
            prev_state.context.pop(0)
            cc = " ".join([paragraph[i] for i in prev_state.context])
            e = create_cpc_input_from_text(
                tokenizer, cc, paragraph[1:], [], group_size=model_size)
          input_map = {
              placeholders.input_ids: e.tokens,
              placeholders.input_mask: e.mask,
              placeholders.segment_ids: e.seg_ids,
              placeholders.softmax_mask: prev_state.softmax_mask
          }
          results = sess.run([model.logits], feed_dict=input_map)
          new_cands = np.argsort(results[0][0])[::-1][:BEAM_SIZE]
          for cand in new_cands:
            tmp_output_order = list(prev_state.output_order)
            tmp_softmax_mask = list(prev_state.softmax_mask)
            tmp_context = list(prev_state.context)
            tmp_output_order.append(cand)
            tmp_softmax_mask[cand] = False
            tmp_context.append(cand + 1)
            next_beam_states.append(
                State(prev_state.prob + results[0][0][cand], tmp_output_order,
                      tmp_softmax_mask, tmp_context))
        # purge states
        beam_states = sorted(
            next_beam_states, key=lambda x: x.prob, reverse=True)[:BEAM_SIZE]
        next_beam_states = []
      output_order = sorted(
          beam_states, key=lambda x: x.prob, reverse=True)[0].output_order
      # append best
      print("Completed %d / %d" % (para_idx, len(paragraphs)))
      tf.logging.info("Completed %d / %d" % (para_idx, len(paragraphs)))
      output_order_list.append(output_order)
  return output_order_list
Compute the predicted output for each paragraph.
language/conpono/evals/coherence_eval.py