code: def create_model(
      bert_config, is_training, input_ids, input_mask, segment_ids
  ):
      """Creates a classification model."""
      model = bert_lib.BertModel(
          config=bert_config,
          is_training=is_training,
          input_ids=input_ids,
          input_mask=input_mask,
          token_type_ids=segment_ids,
          mesh_shape=FLAGS.mesh_shape,
          layout=FLAGS.layout,
      )
  
      final_hidden = model.get_sequence_output()
  
      model_dim = model.model_dim
      class_dim = mtf.Dimension("class", 2)
  
      logits = mtf.layers.dense(
          final_hidden,
          reduced_dims=[model_dim],
          new_dims=[class_dim],
          kernel_initializer=tf.truncated_normal_initializer(
              stddev=0.02
          ),
          name="cls/squad/output",
      )
  
      unstacked_logits = mtf.unstack(logits, dim=class_dim)
  
      (start_logits, end_logits) = (
          unstacked_logits[0],
          unstacked_logits[1],
      )
  
      return (start_logits, end_logits)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import json
  import math
  import os
  import random
  import mesh_tensorflow
  import mesh_tensorflow.bert.bert
  import mesh_tensorflow.bert.optimization
  import mesh_tensorflow.bert.tokenization
  import six
  from six.moves import range
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 5 for calling a global function with example usage: create_model(bert_config, is_training, input_ids, input_mask, segment_ids) and returns: start_logits,  end_logits