code: def make_bert_squad_model_fn(optimizer):
      def _bert_squad_model_fn(features, labels, mode, params):
          import tensorflow as tf
          from zoo.tfpark import ZooOptimizer
  
          final_hidden = bert_model(
              features, labels, mode, params
          ).get_sequence_output()
          final_hidden_shape = modeling.get_shape_list(
              final_hidden, expected_rank=3
          )
          batch_size = final_hidden_shape[0]
          seq_length = final_hidden_shape[1]
          hidden_size = final_hidden_shape[2]
  
          output_weights = tf.get_variable(
              "cls/squad/output_weights",
              [2, hidden_size],
              initializer=tf.truncated_normal_initializer(stddev=0.02),
          )
          output_bias = tf.get_variable(
              "cls/squad/output_bias",
              [2],
              initializer=tf.zeros_initializer(),
          )
  
          final_hidden_matrix = tf.reshape(
              final_hidden, [batch_size * seq_length, hidden_size]
          )
          logits = tf.matmul(
              final_hidden_matrix, output_weights, transpose_b=True
          )
          logits = tf.nn.bias_add(logits, output_bias)
  
          logits = tf.reshape(logits, [batch_size, seq_length, 2])
          logits = tf.transpose(logits, [2, 0, 1])
          unstacked_logits = tf.unstack(logits, axis=0)
          (start_logits, end_logits) = (
              unstacked_logits[0],
              unstacked_logits[1],
          )
  
          if mode == tf.estimator.ModeKeys.TRAIN:
  
              def compute_loss(logits, positions):
                  one_hot_positions = tf.one_hot(
                      positions, depth=seq_length, dtype=tf.float32
                  )
                  log_probs = tf.nn.log_softmax(logits, axis=-1)
                  loss = -tf.reduce_mean(
                      tf.reduce_sum(
                          one_hot_positions * log_probs, axis=-1
                      )
                  )
                  return loss
  
              start_positions = labels["start_positions"]
              end_positions = labels["end_positions"]
  
              start_loss = compute_loss(start_logits, start_positions)
              end_loss = compute_loss(end_logits, end_positions)
              total_loss = (start_loss + end_loss) / 2.0
              train_op = ZooOptimizer(optimizer).minimize(total_loss)
              return tf.estimator.EstimatorSpec(
                  mode=mode, train_op=train_op, loss=total_loss
              )
          elif mode == tf.estimator.ModeKeys.PREDICT:
              predictions = {
                  "unique_ids": features["unique_ids"],
                  "start_logits": start_logits,
                  "end_logits": end_logits,
              }
              return tf.estimator.EstimatorSpec(
                  mode=mode, predictions=predictions
              )
          else:
              raise ValueError(
                  "Currently only TRAIN and PREDICT modes are supported. "
                  "SQuAD uses a separate script for EVAL"
              )
  
      return _bert_squad_model_fn
  
imports: from zoo.tfpark.text.estimator import *
label_desc: Use image node_id 1 for calling a global function with example usage: make_bert_squad_model_fn(optimizer) and returns: _bert_squad_model_fn, tf, loss, tf