Bert: BERT feature extraction for tensorflow serving

[mvss80]

I would like to export BERT feature extraction for use via tensorflow serving.

Following what has been done here: https://github.com/bigboNed3/bert_serving, I modified extract_features.py by adding this function

def serving_input_fn():
    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
    input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids')
    input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask')
    segment_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], 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

and the following code snippet to the main function.

  if FLAGS.do_export:
    estimator._export_to_tpu = False
    estimator.export_savedmodel(FLAGS.export_dir, serving_input_fn, strip_default_attrs=True)

I then run the following on command line

python extract_features.py --input_file=/tmp/input.txt \
  --output_file=/tmp/output.jsonl \
  --vocab_file=$BERT_BASE_DIR/vocab.txt \
  --bert_config_file=$BERT_BASE_DIR/bert_config.json \
  --init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
  --layers=-1,-2,-3,-4 --max_seq_length=128 --batch_size=8 \
  --do_export=True --export_dir='/models'

I get these errors below

Traceback (most recent call last):
  File "extract_features_export.py", line 431, in <module>
    tf.app.run()
  File "/home/m/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/tensorflow/python/platform/app.py", line 125, in run
    _sys.exit(main(argv))
  File "extract_features_export.py", line 423, in main
    estimator.export_savedmodel(FLAGS.export_dir, serving_input_fn, strip_default_attrs=True)
  File "/home/m/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 663, in export_savedmodel
    mode=model_fn_lib.ModeKeys.PREDICT)
  File "/home/m/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 789, in _export_saved_model_for_mode
    strip_default_attrs=strip_default_attrs)
  File "/home/m/anaconda3/envs/gpu-dev/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 878, in _export_all_saved_models
    raise ValueError("Couldn't find trained model at %s." % self._model_dir)
ValueError: Couldn't find trained model at /tmp/tmp3eh_rldt.

Specifying the directory explicitly by updating run_config to include an output_dir

  run_config = tf.contrib.tpu.RunConfig(
      master=FLAGS.master,
      model_dir=FLAGS.output_dir,
      tpu_config=tf.contrib.tpu.TPUConfig(
          num_shards=FLAGS.num_tpu_cores,
          per_host_input_for_training=is_per_host))

and calling the function by giving --output_dir=$BERT_BASE_DIR still gives me the same ValueError: Couldn't find trained model at ....

Has anyone else faced this issue? How can I fix this?

[NatLee]

the same problem : (

[apurvaasf]

I had faced the same problem while exporting a fine-tuned model. Following are the things that I understood:

1. So, your trained model is saved in the system temporary directory and deleted/cleaned by the system when your program was completed.
2. For you to be able to export a saved model, you need to train the model again. Tesnorflow will use the temporary directories to generate a saved model.

do_train and do_export need to be true at the same time in this use case.

I did that and I stopped getting this error. I got the output_dir that I wanted.
Here is my reference.
Hope this helps.

[mvss80]

I had faced the same problem while exporting a fine-tuned model. Following are the things that I understood:

1. So, your trained model is saved in the system temporary directory and deleted/cleaned by the system when your program was completed.
2. For you to be able to export a saved model, you need to train the model again. Tesnorflow will use the temporary directories to generate a saved model.

do_train and do_export need to be true at the same time in this use case.

I did that and I stopped getting this error. I got the output_dir that I wanted.
Here is my reference.
Hope this helps.

Thanks! Can you share the code you used to export?

[stevewyl]

I'm also trying to export the unfine-tuned BERT model as an online service.
I followed the official instructions SavedModel and successfully exported the fine-tuned model.
But when I try to export the original BERT model, it failed.
I compare two model directory, the fine-tuned model dir has two more files: graph.pbtxt and checkpoint. I'm not sure whether this is the reason why unfine-tuned model export failed.

[stevewyl]

@apurvaasf Hi! Is there any possibility that just export the original bert model as SavedModel? Or can we do a fake training process to generate 'checkpoint' file, which is missed in the official model dir, and then export model?

[apurvaasf]

@stevewyl from what I can see, there is a bert_model.ckpt in each of the BERT pre-trained .zip files!

[stevewyl]

@apurvaasf I found an easiest way to export original bert model to SavedModel.

# load the checkpoint from bert
# create an estimator which contains original bert model and some other ops if you want
estimator = ...
save_hook = tf.train.CheckpointSaverHook(FLAGS.output_dir, save_secs=1)
estimator.predict(input_fn=predict_input_fn, hooks=[save_hook])
# now you will get graph.pbtxt which is used in SavedModel, and then
estimator.export_savedmodel(...)

[xielangtc]

any improvement ? I am following https://github.com/bigboNed3/bert_serving. model
loading is ok. but , failed in serving: err msg is " session run failed: Invalid argument: Expects arg[0] to be string but int64 is provided"

[nlp4whp]

@apurvaasf I found an easiest way to export original bert model to SavedModel.

# load the checkpoint from bert
# create an estimator which contains original bert model and some other ops if you want
estimator = ...
save_hook = tf.train.CheckpointSaverHook(FLAGS.output_dir, save_secs=1)
estimator.predict(input_fn=predict_input_fn, hooks=[save_hook])
# now you will get graph.pbtxt which is used in SavedModel, and then
estimator.export_savedmodel(...)

yeah I'm with you, estimator has to run predict at least one time before export_savedmodel.

[fengzuo97]

I write the fake train code, and solved this problem

`

loss = model.get_pooled_output()[0][0]
optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
if mode == tf.estimator.ModeKeys.TRAIN:
    output_spec = tf.contrib.tpu.TPUEstimatorSpec(
    mode=mode, loss=loss, train_op=train_op, scaffold_fn=scaffold_fn
)

`

and add this in the main function

estimator.train(input_fn)

[jiqiujia]

@apurvaasf I found an easiest way to export original bert model to SavedModel.

# load the checkpoint from bert
# create an estimator which contains original bert model and some other ops if you want
estimator = ...
save_hook = tf.train.CheckpointSaverHook(FLAGS.output_dir, save_secs=1)
estimator.predict(input_fn=predict_input_fn, hooks=[save_hook])
# now you will get graph.pbtxt which is used in SavedModel, and then
estimator.export_savedmodel(...)

I add this but I still can't get the same error. It seems that the model doesn't run the predict at all.

[jageshmaharjan]

Maybe check this out if you are looking for serving BERT fine-tuned model.
BERT Serving and Inferencing from fine-tuned

[pqchen]

I had the same error even with graph.pbtxt present in the model folder, which made me believe that perhaps tf estimator was looking into the wrong folder. Turned out the issue can be fixed by providing the correct "checkpoint_path" to the export_saved_model() function.

For example:
estimator.export_saved_model(..., checkpoint_path='model.ckpt-12345')

And this works even without graph.pbtxt.