Saving And Doing Inference With Tensorflow Bert Model

I have created a binary classifier with Tensorflow BERT language model. Here is the link. After the model is trained, it saves the model and produces the following files. Predict

Solution 1:

The create_model function present in notebook takes some arguments. These features are passed to the model.

By updating the serving_input_fn function to following, the serving function works as expected.

Updated Code

def serving_input_receiver_fn():
  feature_spec = {
      "input_ids" : tf.FixedLenFeature([MAX_SEQ_LENGTH], tf.int64),
      "input_mask" : tf.FixedLenFeature([MAX_SEQ_LENGTH], tf.int64),
      "segment_ids" : tf.FixedLenFeature([MAX_SEQ_LENGTH], tf.int64),
      "label_ids" :  tf.FixedLenFeature([], tf.int64)
  }
  serialized_tf_example = tf.placeholder(dtype=tf.string,
                                         shape=[None],
                                         name='input_example_tensor')
  print(serialized_tf_example.shape)
  receiver_tensors = {'example': serialized_tf_example}
  features = tf.parse_example(serialized_tf_example, feature_spec)
  return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

export_path = '/content/drive/My Drive/binary_class/bert/'
estimator._export_to_tpu = False  # this is important
estimator.export_saved_model(export_dir_base=export_path,serving_input_receiver_fn=serving_input_receiver_fn)

Solution 2:

Using the feature_spec dict in the serving_input_receiver_fn did not work for me. I used the serving_input_fn below from someone else asking the same question.

To use the loaded estimator:

def serving_input_fn():
    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
    input_ids = tf.placeholder(tf.int32, [None, MAX_SEQ_LEN], name='input_ids')
    input_mask = tf.placeholder(tf.int32, [None, MAX_SEQ_LEN], name='input_mask')
    segment_ids = tf.placeholder(tf.int32, [None, MAX_SEQ_LEN], name='segment_ids')
    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(
        {
        'label_ids': label_ids,
        'input_ids': input_ids,
        'input_mask': input_mask,
        'segment_ids': segment_ids
        }
    )()
    return input_fn


export_path = '../testing'
estimator._export_to_tpu = False  # this is important
estimator.export_saved_model(export_dir_base=export_path,serving_input_receiver_fn=serving_input_fn)

from tensorflow.contrib import predictor
predict_fn = predictor.from_saved_model('../testing/1589420991')

def predict(sentences, predict_fn):
    labels = [0, 1]
    input_examples = [
        run_classifier.InputExample(
            guid="",
            text_a = x,
            text_b = None,
            label = 0
        ) for x in sentences] # here, "" is just a dummy label
    input_features = run_classifier.convert_examples_to_features(
        input_examples, labels, MAX_SEQ_LEN, tokenizer
    )

    all_input_ids = []
    all_input_mask = []
    all_segment_ids = []
    all_label_ids = []

    for feature in input_features:
        all_input_ids.append(feature.input_ids)
        all_input_mask.append(feature.input_mask)
        all_segment_ids.append(feature.segment_ids)
        all_label_ids.append(feature.label_id)
    pred_dict = {
        'input_ids': all_input_ids,
        'input_mask': all_input_mask,
        'segment_ids': all_segment_ids,
        'label_ids': all_label_ids
    }

    predictions = predict_fn(pred_dict)
    return [
        (sentence, prediction, label)
        for sentence, prediction, label in zip(pred_sentences, predictions['probabilities'], predictions['labels'])
    ]

pred_sentences = [
  "That movie was absolutely awful",
  "The acting was a bit lacking",
  "The film was creative and surprising",
  "Absolutely fantastic!",
]

predictions = predict(pred_sentences, predict_fn)
print(predictions)
[('That movie was absolutely awful',
  array([-0.26713806, -1.4505868 ], dtype=float32),
  0),
 ('The acting was a bit lacking',
  array([-0.23832974, -1.5508994 ], dtype=float32),
  0),
 ('The film was creative and surprising',
  array([-0.2784096, -1.4146391], dtype=float32),
  0),
 ('Absolutely fantastic!',
  array([-0.29031944, -1.3784236 ], dtype=float32),
  0),
 ('The patient has diabetes',
  array([-0.33836085, -1.2480571 ], dtype=float32),
  0),
 ('The patient does not have diabetes',
  array([-0.29378486, -1.3682064 ], dtype=float32),
  0)]