From 57fd73c541589ac37dbedb2e5d393dc6791031aa Mon Sep 17 00:00:00 2001
From: Ivan Zhang <ivan@ivanzhang.ca>
Date: Thu, 6 Jun 2019 12:07:18 -0400
Subject: [PATCH 1/2] return transformed samples in predict, add json print
 option

---
 cli/cmd/predict.go          | 19 +++++++++++++++++--
 pkg/lib/json/json.go        |  9 +++++++++
 pkg/workloads/tf_api/api.py |  1 +
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/cli/cmd/predict.go b/cli/cmd/predict.go
index af7b55b5d9..bf70164480 100644
--- a/cli/cmd/predict.go
+++ b/cli/cmd/predict.go
@@ -33,26 +33,31 @@ import (
 	"github.com/cortexlabs/cortex/pkg/operator/api/resource"
 )
 
+var predictPrintJSON bool
+
 func init() {
+	predictCmd.PersistentFlags().BoolVarP(&predictPrintJSON, "json", "j", false, "print the raw json response")
 	addAppNameFlag(predictCmd)
 	addEnvFlag(predictCmd)
 }
 
 type PredictResponse struct {
 	ResourceID                string                     `json:"resource_id"`
-	ClassificationPredictions []ClassificationPrediction `json:"classification_predictions"`
-	RegressionPredictions     []RegressionPrediction     `json:"regression_predictions"`
+	ClassificationPredictions []ClassificationPrediction `json:"classification_predictions,omitempty"`
+	RegressionPredictions     []RegressionPrediction     `json:"regression_predictions,omitempty"`
 }
 
 type ClassificationPrediction struct {
 	PredictedClass         int         `json:"predicted_class"`
 	PredictedClassReversed interface{} `json:"predicted_class_reversed"`
 	Probabilities          []float64   `json:"probabilities"`
+	TransformedSample      interface{} `json:"transformed_sample"`
 }
 
 type RegressionPrediction struct {
 	PredictedValue         float64     `json:"predicted_value"`
 	PredictedValueReversed interface{} `json:"predicted_value_reversed"`
+	TransformedSample      interface{} `json:"transformed_sample"`
 }
 
 var predictCmd = &cobra.Command{
@@ -87,6 +92,16 @@ var predictCmd = &cobra.Command{
 			errors.Exit(err)
 		}
 
+		if predictPrintJSON {
+			prettyResp, err := json.Pretty(predictResponse)
+			if err != nil {
+				errors.Exit(err)
+			}
+
+			fmt.Println(prettyResp)
+			return
+		}
+
 		apiID := predictResponse.ResourceID
 		api := resourcesRes.APIStatuses[apiID]
 
diff --git a/pkg/lib/json/json.go b/pkg/lib/json/json.go
index d9ccc3327b..c7ed13a493 100644
--- a/pkg/lib/json/json.go
+++ b/pkg/lib/json/json.go
@@ -73,3 +73,12 @@ func WriteJSON(obj interface{}, outPath string) error {
 	}
 	return nil
 }
+
+func Pretty(obj interface{}) (string, error) {
+	b, err := json.MarshalIndent(obj, "", "  ")
+	if err != nil {
+		return "", err
+	}
+
+	return string(b), nil
+}
diff --git a/pkg/workloads/tf_api/api.py b/pkg/workloads/tf_api/api.py
index e57f6be191..4d3e6d69c6 100644
--- a/pkg/workloads/tf_api/api.py
+++ b/pkg/workloads/tf_api/api.py
@@ -193,6 +193,7 @@ def run_predict(sample):
     prediction_request = create_prediction_request(transformed_sample)
     response_proto = local_cache["stub"].Predict(prediction_request, timeout=10.0)
     result = parse_response_proto(response_proto)
+    result["transformed_sample"] = transformed_sample
     util.log_indent("Raw sample:", indent=4)
     util.log_pretty(sample, indent=6)
     util.log_indent("Transformed sample:", indent=4)

From 04cec6cc8dda137276e8211090a685a0f6007562 Mon Sep 17 00:00:00 2001
From: Ivan Zhang <ivan@ivanzhang.ca>
Date: Thu, 6 Jun 2019 12:15:03 -0400
Subject: [PATCH 2/2] update docs

---
 docs/operator/cli.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/operator/cli.md b/docs/operator/cli.md
index 54c14ef9c4..59344ce007 100644
--- a/docs/operator/cli.md
+++ b/docs/operator/cli.md
@@ -58,6 +58,7 @@ Flags:
   -a, --app string   app name
   -e, --env string   environment (default "dev")
   -h, --help         help for predict
+  -j, --json         print the raw json response
 ```
 
 The `predict` command converts samples from a JSON file into prediction requests and outputs the response. This command is useful for quickly testing model output.