def model_function(features, labels, mode, params, embeddings):
  """A model function satisfying the tf.estimator API.

  Args:
    features: Dictionary of feature tensors with keys:
        - question_tok: <string> [batch_size, max_question_len]
        - context_tok: <string> [batch_size, max_num_context, max_context_len]
        - question_tok_len: <int32> [batch_size]
        - num_context: <int32> [batch_size]
        - context_tok_len: <int32> [batch_size]
        - question_tok_wid: <int32> [batch_size, max_question_len]
        - context_tok_wid: <int32> [batch_size, max_num_context,
          max_context_len]
         - long_answer_indices: <int32> [batch_size]
    labels: <int32> [batch_size] for answer index (-1 = NULL).
    mode: One of the keys from tf.estimator.ModeKeys.
    params: Dictionary of hyperparameters.
    embeddings: An embedding_utils.PretrainedWordEmbeddings object.

  Returns:
    estimator_spec: A tf.estimator.EstimatorSpec object.
  """
  del params  # Unused.

  if mode == tf_estimator.ModeKeys.PREDICT:
    # Add a dummy batch dimension if we are exporting the predictor.
    features = {k: tf.expand_dims(v, 0) for k, v in features.items()}

  embedding_weights, embedding_scaffold = embeddings.get_params(trainable=False)

  # Features.
  question_tok_len = features["question_tok_len"]
  question_tok_wid = features["question_tok_wid"]
  context_tok_wid = features["context_tok_wid"]
  num_context = features["num_context"]
  context_tok_len = features["context_tok_len"]

  # Truncate the contexts and labels to a certain maximum length.
  context_tok_wid, num_context, context_tok_len = (
      nq_long_utils.truncate_contexts(
          context_token_ids=context_tok_wid,
          num_contexts=num_context,
          context_len=context_tok_len,
          max_contexts=FLAGS.max_contexts,
          max_context_len=FLAGS.max_context_len))

  non_null_context_scores = nq_long_decatt_model.build_model(
      question_tok_wid=question_tok_wid,
      question_lens=question_tok_len,
      context_tok_wid=context_tok_wid,
      context_lens=context_tok_len,
      embedding_weights=embedding_weights,
      mode=mode)

  # Mask out contexts that are padding.
  num_context_mask = tf.log(
      tf.sequence_mask(
          num_context,
          tensor_utils.shape(non_null_context_scores, 1),
          dtype=tf.float32))
  non_null_context_scores += num_context_mask

  # <float> [batch_size, 1]
  null_score = tf.zeros([tf.shape(question_tok_wid)[0], 1])

  # Offset everything by 1 to account for null context.
  # [batch_size, 1 + max_contexts]
  context_scores = tf.concat([null_score, non_null_context_scores], 1)

  if mode != tf_estimator.ModeKeys.PREDICT:
    labels = nq_long_utils.truncate_labels(labels, FLAGS.max_contexts)

    # In the data, NULL is given index -1 but this is not compatible with
    # softmax so shift by 1.
    labels = labels + 1

    # Reweight null examples.
    weights = nq_long_utils.compute_null_weights(labels, FLAGS.null_weight)

    # When computing the loss we take only the first label.
    loss_labels = labels[:, 0]

    # []
    loss = tf.losses.sparse_softmax_cross_entropy(
        labels=loss_labels, logits=context_scores, weights=weights)

    optimizer = tf.train.AdagradOptimizer(learning_rate=FLAGS.learning_rate)
    train_op = optimizer.minimize(
        loss=loss, global_step=tf.train.get_global_step())

    # <int32> [batch_size]
    eval_predictions = tf.to_int32(tf.argmax(context_scores, 1))

    non_null_match, non_null_gold, non_null_predictions = (
        nq_long_utils.compute_match_stats(eval_predictions, labels))

    precision, precision_op = (
        tf.metrics.mean(non_null_match, weights=non_null_predictions))
    recall, recall_op = (tf.metrics.mean(non_null_match, weights=non_null_gold))

    f1, f1_op = (
        nq_long_utils.f1_metric(
            precision=precision,
            precision_op=precision_op,
            recall=recall,
            recall_op=recall_op))

    # Bogus metric until we figure out how to connect Ming Wei's eval code.
    eval_metric_ops = {
        "precision": (precision, precision_op),
        "recall": (recall, recall_op),
        "f1": (f1, f1_op)
    }
  else:
    loss = None
    train_op = None
    eval_metric_ops = {}

  # In the export, we never predict NULL since the eval metric will compute the
  # best possible F1.
  export_long_answer_idx = tf.to_int32(tf.argmax(non_null_context_scores, 1))
  export_long_answer_score = tf.reduce_max(non_null_context_scores, 1)
  predictions = dict(idx=export_long_answer_idx, score=export_long_answer_score)

  if mode == tf_estimator.ModeKeys.PREDICT:
    # Remove the dummy batch dimension if we are exporting the predictor.
    predictions = {k: tf.squeeze(v, 0) for k, v in predictions.items()}

  estimator_spec = tf_estimator.EstimatorSpec(
      mode=mode,
      loss=loss,
      predictions=predictions,
      train_op=train_op,
      eval_metric_ops=eval_metric_ops,
      scaffold=embedding_scaffold)

  return estimator_spec
A model function satisfying the tf.estimator API.

  Args:
    features: Dictionary of feature tensors with keys:
        - question_tok: <string> [batch_size, max_question_len]
        - context_tok: <string> [batch_size, max_num_context, max_context_len]
        - question_tok_len: <int32> [batch_size]
        - num_context: <int32> [batch_size]
        - context_tok_len: <int32> [batch_size]
        - question_tok_wid: <int32> [batch_size, max_question_len]
        - context_tok_wid: <int32> [batch_size, max_num_context,
          max_context_len]
         - long_answer_indices: <int32> [batch_size]
    labels: <int32> [batch_size] for answer index (-1 = NULL).
    mode: One of the keys from tf.estimator.ModeKeys.
    params: Dictionary of hyperparameters.
    embeddings: An embedding_utils.PretrainedWordEmbeddings object.

  Returns:
    estimator_spec: A tf.estimator.EstimatorSpec object.
  
language/question_answering/decatt_docreader/models/nq_long_model.py