Respond with transformed samples in prediction requests, add json print option (#153)

Author: 1vn

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]
 

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.

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
+}

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)