Update sentiment.py

Author: ospillinger

examples/sentiment/sentiment.py

@@ -2,15 +2,16 @@
 import tensorflow_hub as hub
 from bert import tokenization, run_classifier
 
+labels = ["negative", "positive"]
+
 with tf.Graph().as_default():
     bert_module = hub.Module("https://tfhub.dev/google/bert_uncased_L-12_H-768_A-12/1")
     tokenization_info = bert_module(signature="tokenization_info", as_dict=True)
-    vocab_file = tokenization_info["vocab_file"]
-    do_lower_case = tokenization_info["do_lower_case"]
     with tf.Session() as sess:
-        vocab_file, do_lower_case = sess.run([vocab_file, do_lower_case])
-
-tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=do_lower_case)
+        vocab_file, do_lower_case = sess.run(
+            [tokenization_info["vocab_file"], tokenization_info["do_lower_case"]]
+        )
+        tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=do_lower_case)
 
 
 def pre_inference(sample, metadata):
@@ -20,5 +21,4 @@ def pre_inference(sample, metadata):
 
 
 def post_inference(prediction, metadata):
-    labels = ["negative", "positive"]
-    return {"sentiment": labels[prediction["response"]["labels"][0]]}
+    return labels[prediction["response"]["labels"][0]]