//  [python]
def create_model(model,
                 labels,
                 decoder_inputs,
                 batch_size,
                 model_type="decode",
                 sep_positions=None):
  """Creates a classification model.

  Args:
    model: the BERT model from modeling.py
    labels: ground truth paragraph order
    decoder_inputs: the input to the decoder if used
    batch_size: the batch size
    model_type: one of decode, pooled, attn
    sep_positions: (optional) for "pooled" indecies of SEP tokens

  Returns:
    tuple of (loss, per_example_loss, logits, probabilities) for model
  """

  output_layer = model.get_pooled_output()

  hidden_size = output_layer.shape[-1].value
  tpu_batch_size = tf.shape(output_layer)[0]

  num_labels = 5  # GOOGLE-INTERNAL TODO(daniter) this shouldn't be hardcoded

  with tf.variable_scope("paragraph_reconstruct"):
    if model_type == "decode":
      lstm_cell = tf.nn.rnn_cell.LSTMCell(
          num_units=hidden_size, use_peepholes=True, state_is_tuple=True)

      def sample_fn(x):
        return tf.cast(tf.reshape(tf.argmax(x, axis=-1), (-1, 1)), tf.float32)

      helper = FixedSizeInferenceHelper(
          sample_fn=sample_fn,
          sample_shape=[1],
          sample_dtype=tf.float32,
          start_inputs=decoder_inputs[:, 0],
          end_fn=None)

      # Decoder
      project_layer = tf.layers.Dense(
          num_labels, use_bias=False, name="output_projection")

      my_decoder = contrib_seq2seq.BasicDecoder(
          lstm_cell,
          helper,
          tf.nn.rnn_cell.LSTMStateTuple(output_layer, output_layer),
          output_layer=project_layer)

      # Dynamic decoding
      outputs, _, _ = contrib_seq2seq.dynamic_decode(
          my_decoder,
          swap_memory=True,
          scope="paragraph_reconstruct",
          maximum_iterations=5)

      logits = outputs.rnn_output

      cross_ent = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=labels, logits=logits)

      per_example_loss = cross_ent
      loss = tf.reduce_sum(cross_ent) / tf.cast(batch_size, tf.float32)
      probabilities = tf.nn.softmax(logits, axis=-1)

    # GOOGLE-INTERAL: TODO(daniter)  currently neither of these actually train
    elif model_type == "pooled":
      token_embeddings = model.get_sequence_output()
      # sep positions come out batch by batch so we need to add the batch index
      # we do that explicitly here since we don't know the batch size in the
      # record decoder
      batch_idx = tf.range(tpu_batch_size)
      batch_idx = tf.reshape(batch_idx, [tpu_batch_size, 1])
      batch_idx = tf.tile(batch_idx, [1, 5])  # double check
      batch_idx = tf.reshape(batch_idx, [tpu_batch_size, 5, 1])
      # batch_idx = tf.Print(batch_idx, [batch_idx],
      #                      message="batch_idx", summarize=999999)
      sep_positions = tf.concat([batch_idx, sep_positions], axis=2)
      # sep_positions = tf.Print(sep_positions, [sep_positions],
      #                          message="sep_positions", summarize=999999)

      sep_vecs = tf.gather_nd(token_embeddings, sep_positions)
      sep_vecs = tf.reshape(sep_vecs, [tpu_batch_size, 5, hidden_size])
      # sep_vecs = tf.Print(sep_vecs, [sep_vecs], message="sep_vecs",
      #                     summarize=999999)

      logits = tf.layers.dense(
          inputs=sep_vecs, units=num_labels, name="output_projection")
      # logits = tf.Print(logits, [logits], message="logits", summarize=999999)
      cross_ent = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=labels, logits=logits)

      per_example_loss = cross_ent
      loss = tf.reduce_sum(cross_ent) / tf.cast(batch_size, tf.float32)
      probabilities = tf.nn.softmax(logits, axis=-1)

    elif model_type == "attn":
      # change size to match sequence embedding size
      input_consts = tf.constant([0, 1, 2, 3, 4])
      position_encoding = tf.broadcast_to(input_consts, [tpu_batch_size, 5])
      # position_encoding = tf.to_float(
      # tf.reshape(position_encoding, (-1, 5, 1)))
      token_type_table = tf.get_variable(
          name="attention_embedding",
          shape=[5, 512],  # don't hardcode
          initializer=tf.truncated_normal_initializer(stddev=0.02))
      # This vocab will be small so we always do one-hot here, since it is
      # always faster for a small vocabulary.
      flat_token_type_ids = tf.reshape(position_encoding, [-1])
      one_hot_ids = tf.one_hot(flat_token_type_ids, depth=5)
      token_type_embeddings = tf.matmul(one_hot_ids, token_type_table)
      token_type_embeddings = tf.reshape(token_type_embeddings,
                                         [tpu_batch_size, 5, 512])

      token_embeddings = model.get_sequence_output()
      attn = modeling.attention_layer(token_type_embeddings, token_embeddings)
      attn = tf.reshape(attn, (-1, 5, 512))  # head size
      logits = tf.layers.dense(
          inputs=attn, units=num_labels, name="output_projection")
      cross_ent = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=labels, logits=logits)

      per_example_loss = cross_ent
      loss = tf.reduce_sum(cross_ent) / tf.cast(batch_size, tf.float32)
      probabilities = tf.nn.softmax(logits, axis=-1)

  return (loss, per_example_loss, logits, probabilities)