Fix review comments

elibixby · elibixby · commit 8c3da38b2a1b · 2017-02-27T17:38:48.000-08:00
diff --git a/ml_engine/online_prediction/README.md b/ml_engine/online_prediction/README.md
@@ -1,38 +1 @@
-# Online Prediction with the Cloud Machine Learning Engine
-
-This sample assumes that you have already run the Cloud ML Engine end to end walkthrough using the UCI Census dataset, including training a wide & deep model and deploying it into production, ready to take prediction requests.
-
-This sample then shows you how to create a client that takes user JSON inputs, and sends those inputs as online prediction requests to a given deployed model.
-
-In order to use this client, first obtain the following information, and store it into the given environment variables:
-
-```
-PROJECT=<your project name>
-MODEL=<your model name>
-VERSION=<version of model you are using>
-```
-
-Next, launch this client as follows:
-
-```
-python predict.py --project=$PROJECT --model=$MODEL --version=$VERSION
-```
-
-After having done that, the client will ask you for ‘Valid JSON’ input as follows:
-
-```
-Valid JSON >>>
-```
-
-Now you can input a JSON example that corresponds to the schema of the given model. For instance if you are sending prediction requests to the census-based model created in the Cloud ML Engine walkthrough, you can send a JSON example like the following:
-
-```
-{"age": 25, "workclass": " Private", "education": " 11th", "education_num": 7, "marital_status": " Never-married", "occupation": " Machine-op-inspct", "relationship": " Own-child", "race": " Black", "gender": " Male", "capital_gain": 0, "capital_loss": 0, "hours_per_week": 40, "native_country": " United-States"}
-```
-The result should be something along the following lines (depending on how you trained the model/what parameters you used, the results may vary):
-
-```
-[{u'probabilities': [0.992774486541748, 0.007225471083074808], u'logits': [-4.922891139984131], u'classes': 0, u'logistic': [0.007225471083074808]}]
-```
-
-Now that you have a working client, you can adapt this to your own use cases!
+https://cloud.google.com/ml-engine/docs/concepts/prediction-overview
diff --git a/ml_engine/online_prediction/predict.py b/ml_engine/online_prediction/predict.py
@@ -11,8 +11,6 @@
 # the License.
 
 """Examples of using the Cloud ML Engine's online prediction service."""
-from __future__ import print_function
-
 import argparse
 import base64
 import json
@@ -21,6 +19,8 @@
 import googleapiclient.discovery
 # [END import_libraries]
 
+import six
+
 
 # [START predict_json]
 def predict_json(project, model, instances, version=None):
@@ -29,9 +29,10 @@ def predict_json(project, model, instances, version=None):
     Args:
         project (str): project where the Cloud ML Engine Model is deployed.
         model (str): model name.
-        instances ([Mapping[str: any]]): dictionaries from string keys
-            defined by the model deployment, to data with types that match
-            expected tensors
+        instances ([Mapping[str: Any]]): Keys should be the names of Tensors
+            your deployed model expects as inputs. Values should be datatypes
+            convertible to Tensors, or (potentially nested) lists of datatypes
+            convertible to tensors.
         version: str, version of the model to target.
     Returns:
         Mapping[str: any]: dictionary of prediction results defined by the
@@ -106,20 +107,26 @@ def census_to_example_bytes(json_instance):
     for details.
 
     Args:
-        json_instance (Mapping[str: any]): representing data to be serialized.
+        json_instance (Mapping[str: Any]): Keys should be the names of Tensors
+            your deployed model expects to parse using it's tf.FeatureSpec.
+            Values should be datatypes convertible to Tensors, or (potentially
+            nested) lists of datatypes convertible to tensors.
     Returns:
         str: A string as a container for the serialized bytes of
             tf.train.Example protocol buffer.
     """
     import tensorflow as tf
     feature_dict = {}
     for key, data in json_instance.iteritems():
-        if isinstance(data, str) or isinstance(data, unicode):
+        if isinstance(data, six.string_types):
             feature_dict[key] = tf.train.Feature(
                 bytes_list=tf.train.BytesList(value=[str(data)]))
-        elif isinstance(data, int) or isinstance(data, float):
+        elif isinstance(data, float):
             feature_dict[key] = tf.train.Feature(
                 float_list=tf.train.FloatList(value=[data]))
+        elif isinstance(data, int):
+            feature_dict[key] = tf.train.Feature(
+                int64_list=tf.train.Int64List(value=[data]))
     return tf.train.Example(
         features=tf.train.Features(
             feature=feature_dict
@@ -181,4 +188,9 @@ def main(project, model, version=None, force_tfrecord=False):
         default=False
     )
     args = parser.parse_args()
-    main(**args.__dict__)
+    main(
+        args.project,
+        args.model,
+        version=args.version,
+        force_tfrecord=args.force_tfrecord
+    )
diff --git a/ml_engine/online_prediction/predict_test.py b/ml_engine/online_prediction/predict_test.py
@@ -8,12 +8,14 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations under
 # the License.
+
 """Tests for predict.py ."""
+
 import base64
 
 import pytest
 
-from predict import census_to_example_bytes, predict_json
+import predict
 
 
 MODEL = 'census'
@@ -43,26 +45,21 @@
 
 
 def test_predict_json():
-    result = predict_json(PROJECT, MODEL, [JSON, JSON], version=VERSION)
+    result = predict.predict_json(
+        PROJECT, MODEL, [JSON, JSON], version=VERSION)
     assert [EXPECTED_OUTPUT, EXPECTED_OUTPUT] == result
 
 
 def test_predict_json_error():
     with pytest.raises(RuntimeError):
-        predict_json(PROJECT, MODEL, [{"foo": "bar"}], version=VERSION)
+        predict.predict_json(PROJECT, MODEL, [{"foo": "bar"}], version=VERSION)
 
 
-# TODO(elibixby) Run on Travis when TensorFlow PyPi package supports
-# Ubuntu 12.04 See:
-# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/g3doc/get_started/os_setup.md#import-error
 @pytest.mark.slow
 def test_census_example_to_bytes():
-    b = census_to_example_bytes(JSON)
+    b = predict.census_to_example_bytes(JSON)
     assert base64.b64encode(b) is not None
 
 
 def test_predict_tfrecord():
-    # Using the same model for TFRecords and
-    # JSON is currently broken.
-    # TODO(elibixby) when b/35742966 is fixed add
     pass
diff --git a/ml_engine/online_prediction/requirements.txt b/ml_engine/online_prediction/requirements.txt
@@ -1 +1 @@
-tensorflow>=1.0.0
+tensorflow==1.0.0
