tfx_bsl.public.beam.run_inference.RunInferencePerModel

Vectorized variant of RunInference (useful for ensembles).

examples A PCollection containing examples of the following possible kinds, each with their corresponding return type.

  • PCollection[Example] -> PCollection[ Tuple[PredictionLog, ...]]

    • Works with Classify, Regress, MultiInference, Predict and RemotePredict.
  • PCollection[SequenceExample] -> PCollection[ Tuple[PredictionLog, ...]]

    • Works with Predict and (serialized) RemotePredict.
  • PCollection[bytes] -> PCollection[ Tuple[PredictionLog, ...]]

    • For serialized Example: Works with Classify, Regress, MultiInference, Predict and RemotePredict.
    • For everything else: Works with Predict and RemotePredict.
  • PCollection[Tuple[K, Example]] -> PCollection[ Tuple[K, Tuple[PredictionLog, ...]]]

    • Works with Classify, Regress, MultiInference, Predict and RemotePredict.
  • PCollection[Tuple[K, SequenceExample]] -> PCollection[ Tuple[K, Tuple[PredictionLog, ...]]]

    • Works with Predict and (serialized) RemotePredict.
  • PCollection[Tuple[K, bytes]] -> PCollection[ Tuple[K, Tuple[PredictionLog, ...]]]

    • For serialized Example: Works with Classify, Regress, MultiInference, Predict and RemotePredict.
    • For everything else: Works with Predict and RemotePredict.
inference_spec_types A flat iterable of Model inference endpoints. Inference will happen in a fused fashion (ie without data materialization), sequentially across Models within a Beam thread (but in parallel across threads and workers).

A PCollection (possibly keyed) containing a Tuple of prediction logs. The Tuple of prediction logs is 1-1 aligned with inference_spec_types.