The module in question is kubeflow_batch_predict.dataflow.batch_prediction.py
and the DoFn is PredictionDoFn
.
This issue highlights a shortcoming of the current state and will serve as the primary discussion venue for the problem.
In its current state, this DoFn accepts a serialized JSON containing the following format for the examples:
{'instances': [ {'input': TFRecord}, ... ] }
where each item in the list is a dictionary containing the 'input' key and a TFRecord (could be base64 encoding of the TFRecord). This input does not allow for any extraneous top level keys and only yields back a list of inputs and outputs.
The need for extraneous keys exists because there might be extra metadata along with each element which might be needed to identify the input. For instance, consider a prediction task to embed "movies" as high-dimensional vectors. If we desire to write the final results into a CSV file, we would want each row to have extra metadata like "name" etc and we might want this to be passed around in a dict in the DataFlow pipeline (as PCollections
).
This is not possible currently because this would violate the input format to PredictionDoFn
and instead we would have to morph these values into something acceptable. This step is expected however any downstream DoFns that derive PCollections
from PredictionDoFn
will not be able to access any pre-transformed data. Instead all we will have is a list of high-dimensional vectors with no way to relate back to the actual human readable information like "name".
We need to come to a design spec surrounding this so as to accommodate the most generic use cases around Batch Prediction