def model_function(features, labels, mode, vocab, hps):
  """Model function.

  Atttention seq2seq model, augmented with an encoder
  over the targets of the nearest neighbors.

  Args:
    features: Dictionary of input Tensors.
    labels: Unused.
    mode: train or eval. Keys from tf.estimator.ModeKeys.
    vocab: A list of strings of words in the vocabulary.
    hps: Hyperparams.

  Returns:
    A tf.estimator.EstimatorSpec object containing model outputs.
  """
  targets = features["targets"]
  del labels

  if hps.model in ["hypernet"]:
    model = hypernet
  elif hps.model in ["seq2seq", "nn2seq"]:
    model = baseline

  # [batch_size, dec_len]
  decoder_inputs = features["decoder_inputs"]
  encoder_outputs = model.encoder(
      features=features,
      mode=mode,
      vocab=vocab,
      hps=hps
  )

  if mode == tf_estimator.ModeKeys.TRAIN:
    if hps.predict_mode:
      sys.exit(0)
    outputs, _ = model.basic_decoder(
        features=features,
        mode=mode,
        vocab=vocab,
        encoder_outputs=encoder_outputs,
        hps=hps
    )
    decoder_outputs = outputs.decoder_outputs
    decoder_len = features["decoder_len"]

    weights = tf.sequence_mask(
        decoder_len,
        tf.shape(decoder_inputs)[1],
        dtype=tf.float32)
    loss, train_op = common.optimize_log_loss(
        decoder_tgt=targets,
        decoder_outputs=decoder_outputs,
        weights=weights,
        hps=hps)

    return tf_estimator.EstimatorSpec(
        mode=mode,
        loss=loss,
        train_op=train_op)
  elif mode == tf_estimator.ModeKeys.EVAL:
    outputs, teacher_outputs = model.basic_decoder(
        features=features,
        mode=mode,
        vocab=vocab,
        encoder_outputs=encoder_outputs,
        hps=hps
    )
    if hps.beam_width > 0:
      assert hps.beam_width > 1
      outputs = model.beam_decoder(
          features=features,
          mode=mode,
          vocab=vocab,
          encoder_outputs=encoder_outputs,
          hps=hps
      )

    # [batch_size, max_dec_len]
    weights = tf.sequence_mask(
        features["decoder_len"],
        tf.shape(decoder_inputs)[1],
        dtype=tf.float32)

    loss, _ = common.optimize_log_loss(
        decoder_tgt=targets,
        decoder_outputs=teacher_outputs.decoder_outputs,
        weights=weights,
        hps=hps)

    decoder_outputs = outputs.decoder_outputs
    decoder_len = outputs.decoder_len
    if hps.beam_width == 0:
      predicted_ids = decoder_outputs.sample_id
      lengths = outputs.decoder_len
      predictions = {
          "outputs": predicted_ids,
          "lengths": lengths,
      }
    else:
      predicted_ids, lengths = tf.py_func(
          partial(remove_repetitive_trigram,
                  vocab=vocab,
                  hps=hps),
          [decoder_outputs.predicted_ids, outputs.decoder_len],
          (tf.int32, tf.int32)
      )
      predictions = {
          "outputs": predicted_ids,
          "lengths": lengths,
      }
    reference = {
        "outputs": features["reference"],
        "lengths": features["reference_len"]
    }
    eval_metric_ops = rouge_utils.get_metrics(
        predictions["outputs"], predictions["lengths"],
        reference["outputs"], reference["lengths"],
        vocab=vocab, use_bpe=hps.use_bpe, predict_mode=hps.predict_mode
    )

    # [batch_size, src_len]
    src_inputs = features["src_inputs"]
    src_len = features["src_len"]

    # [batch_size, neighbor_len]
    neighbor_inputs = features["neighbor_inputs"]
    neighbor_len = features["neighbor_len"]

    print_op = common.print_text(
        [("source", src_inputs, src_len, True),
         ("neighbor", neighbor_inputs, neighbor_len, True),
         ("targets", reference["outputs"], reference["lengths"], False),
         ("predictions", predictions["outputs"], predictions["lengths"], True)],
        vocab, use_bpe=hps.use_bpe, predict_mode=hps.predict_mode)
    with tf.control_dependencies([print_op]):
      loss = tf.identity(loss)
    return tf_estimator.EstimatorSpec(
        mode=mode,
        loss=loss,
        predictions=predictions,
        eval_metric_ops=eval_metric_ops)
  elif mode == tf_estimator.ModeKeys.PREDICT:
    outputs, teacher_outputs = model.basic_decoder(
        features=features,
        mode=mode,
        vocab=vocab,
        encoder_outputs=encoder_outputs,
        hps=hps
    )
    if hps.beam_width > 0:
      assert hps.beam_width > 1
      outputs = model.beam_decoder(
          features=features,
          mode=mode,
          vocab=vocab,
          encoder_outputs=encoder_outputs,
          hps=hps
      )

    decoder_outputs = outputs.decoder_outputs
    decoder_len = outputs.decoder_len
    if hps.beam_width == 0:
      predicted_ids = decoder_outputs.sample_id
      lengths = outputs.decoder_len
    else:
      predicted_ids = decoder_outputs.predicted_ids[Ellipsis, 0]
      lengths = outputs.decoder_len[Ellipsis, 0]

    predictions = {
        "outputs": predicted_ids,
        "lengths": lengths - 1,
    }
    return tf_estimator.EstimatorSpec(
        mode=mode,
        predictions=predictions)
Model function.

  Atttention seq2seq model, augmented with an encoder
  over the targets of the nearest neighbors.

  Args:
    features: Dictionary of input Tensors.
    labels: Unused.
    mode: train or eval. Keys from tf.estimator.ModeKeys.
    vocab: A list of strings of words in the vocabulary.
    hps: Hyperparams.

  Returns:
    A tf.estimator.EstimatorSpec object containing model outputs.
  
language/labs/exemplar_decoding/models/model_function.py