Added support for Tensorflow Serving. Breaking change to remove

redundant output dimension.

Author: Alan Patterson

README.md

@@ -123,6 +123,26 @@ close to 2x the speed.
 
 The training script has this option added: `train_classifier.py`.
 
+# Tensorflow Serving
+
+As well as using `predictor.py` to run a saved model to provide
+predictions, it is easy to serve a saved model using Tensorflow
+Serving and gRPC. There is a supplied simple rpc client that provides
+predictions by using a gRPC server.
+
+First make sure you install the tensorflow serving binaries. Instructions are [here](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/setup.md#installing-the-modelserver).
+
+You then serve the latest saved model by supplying the base export
+directory where you exported saved models to. This directory will
+contain the numbered model directories:
+
+    tensorflow_model_server --port=9000 --model_base_path=model
+
+Now you can make requests to the server using gRPC calls. An example
+simple client is provided in `predictor_client.py`:
+
+    predictor_client.py --text="Some text to classify"
+
 # Facebook Examples
 
 << NOT IMPLEMENTED YET >>

classifier.py

@@ -106,7 +106,6 @@ def model_fn(features, labels, mode, params):
             -0.1, 0.1))
         text_embedding = tf.reduce_mean(tf.nn.embedding_lookup(
             text_embedding_w, text_ids), axis=-2)
-        text_embedding = tf.expand_dims(text_embedding, -2)
         input_layer = text_embedding
         if FLAGS.use_ngrams:
             ngram_hash = tf.string_to_hash_bucket(features["ngrams"],
@@ -131,8 +130,6 @@ def model_fn(features, labels, mode, params):
             labels = label_lookup_table.lookup(labels)
             loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
                 labels=labels, logits=logits))
-            # Squeeze dimensions from labels and switch to 0-offset
-            labels = tf.squeeze(labels, -1)
             opt = tf.train.AdamOptimizer(params["learning_rate"])
             if FLAGS.horovod:
                 opt = hvd.DistributedOptimizer(opt)

inputs.py

@@ -15,15 +15,15 @@ def BuildTextExample(text, ngrams=None, label=None):
     if ngrams is not None:
         ngrams = [tf.compat.as_bytes(x) for x in ngrams]
         record.features.feature["ngrams"].bytes_list.value.extend(ngrams)
-    return record.SerializeToString()
+    return record
 
 
 def ParseSpec(use_ngrams, include_target):
     parse_spec = {"text": tf.VarLenFeature(dtype=tf.string)}
     if use_ngrams:
         parse_spec["ngrams"] = tf.VarLenFeature(dtype=tf.string)
     if include_target:
-        parse_spec["label"] = tf.FixedLenFeature(shape=(1,), dtype=tf.string,
+        parse_spec["label"] = tf.FixedLenFeature(shape=(), dtype=tf.string,
                                                  default_value=None)
     return parse_spec
 

predictor.py

@@ -41,6 +41,7 @@ def RunModel(saved_model_dir, signature_def_key, tag, text, ngrams_list=None):
         ngrams_list = text_utils.ParseNgramsOpts(ngrams_list)
         ngrams = text_utils.GenerateNgrams(text, ngrams_list)
     example = inputs.BuildTextExample(text, ngrams=ngrams)
+    example = example.SerializeToString()
     inputs_feed_dict = {
         signature_def.inputs["inputs"].name: [example],
     }
@@ -64,8 +65,8 @@ def main(_):
     outputs = RunModel(FLAGS.saved_model, FLAGS.signature_def, FLAGS.tag,
                        FLAGS.text, FLAGS.ngrams)
     if FLAGS.signature_def == "proba":
-        print("Proba:", outputs[0])
-        print("Class(1-N):", np.argmax(outputs[0]) + 1)
+        print("Proba:", outputs)
+        print("Class(1-N):", np.argmax(outputs) + 1)
     elif FLAGS.signature_def == "embedding":
         print(outputs[0])
 

predictor_client.py

@@ -0,0 +1,61 @@
+"""Predict classification on provided text.
+
+Send request to a tensorflow_model_server.
+
+   tensorflow_model_server --port=9000 --model_base_path=$export_dir_base
+
+Usage:
+   
+   predictor_client.py --text='some text' --ngrams=1,2,4
+
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import inputs
+import text_utils
+
+from grpc.beta import implementations
+from tensorflow_serving.apis import classification_pb2
+from tensorflow_serving.apis import prediction_service_pb2
+
+
+tf.flags.DEFINE_string('server', 'localhost:9000',
+                       'TensorflowService host:port')
+tf.flags.DEFINE_string("text", None, "Text to predict label of")
+tf.flags.DEFINE_string("ngrams", None, "List of ngram lengths, E.g. --ngrams=2,3,4")
+tf.flags.DEFINE_string("signature_def", "proba",
+                       "Stored signature key of method to call (proba|embedding)")
+FLAGS = tf.flags.FLAGS
+
+
+def Request(text, ngrams):
+    text = text_utils.TokenizeText(text)
+    ngrams = None
+    if ngrams is not None:
+        ngrams_list = text_utils.ParseNgramsOpts(ngrams)
+        ngrams = text_utils.GenerateNgrams(text, ngrams_list)
+    example = inputs.BuildTextExample(text, ngrams=ngrams)
+    request = classification_pb2.ClassificationRequest()
+    request.model_spec.name = 'default'
+    request.model_spec.signature_name = 'proba'
+    request.input.example_list.examples.extend([example])
+    return request
+
+
+def main(_):
+    if not FLAGS.text:
+        raise ValueError("No --text provided")
+    host, port = FLAGS.server.split(':')
+    channel = implementations.insecure_channel(host, int(port))
+    stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)
+    request = Request(FLAGS.text, FLAGS.ngrams)
+    result = stub.Classify(request, 10.0)  # 10 secs timeout
+    print(result)
+
+
+if __name__ == '__main__':
+    tf.app.run()
+    

process_input.py

@@ -94,7 +94,7 @@ def WriteExamples(examples, outputfile, num_shards):
                                                  (shard, num_shards))
         record = inputs.BuildTextExample(
             example["text"], example.get("ngrams", None), example["label"])
-        writer.write(record)
+        writer.write(record.SerializeToString())
 
 
 def WriteVocab(examples, vocabfile, labelfile):