code: def make_bert_ner_model_fn(optimizer):
      import tensorflow as tf
      from zoo.tfpark import ZooOptimizer
  
      def _bert_ner_model_fn(features, labels, mode, params):
          output_layer = bert_model(
              features, labels, mode, params
          ).get_sequence_output()
          if mode == tf.estimator.ModeKeys.TRAIN:
              output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)
          logits = tf.layers.dense(output_layer, params["num_entities"])
          if mode == tf.estimator.ModeKeys.TRAIN:
              logits = tf.reshape(logits, [-1, params["num_entities"]])
              labels = tf.reshape(labels, [-1])
              mask = tf.cast(features["input_mask"], dtype=tf.float32)
              one_hot_labels = tf.one_hot(
                  labels, depth=params["num_entities"], dtype=tf.float32
              )
              loss = tf.losses.softmax_cross_entropy(
                  logits=logits, onehot_labels=one_hot_labels
              )
              loss *= tf.reshape(mask, [-1])
              loss = tf.reduce_sum(loss)
              total_size = tf.reduce_sum(mask)
              total_size += (
                  1e-12  # to avoid division by 0 for all-0 weights
              )
              loss /= total_size
              train_op = ZooOptimizer(optimizer).minimize(loss)
              return tf.estimator.EstimatorSpec(
                  mode=mode, train_op=train_op, loss=loss
              )
          elif mode == tf.estimator.ModeKeys.PREDICT:
              probabilities = tf.nn.softmax(logits, axis=-1)
              predict = tf.argmax(probabilities, axis=-1)
              return tf.estimator.EstimatorSpec(
                  mode=mode, predictions=predict
              )
          else:
              raise ValueError(
                  "Currently only TRAIN and PREDICT modes are supported for NER"
              )
  
      return _bert_ner_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_ner_model_fn(optimizer) and returns: _bert_ner_model_fn, tf, tf