tfx_bsl.public.beam.run_inference.RunInferenceOnKeyedBatches

Run inference over pre-batched keyed inputs.

tfx_bsl.public.beam.run_inference.RunInferenceOnKeyedBatches(
    inference_spec_type: tfx_bsl.public.proto.model_spec_pb2.InferenceSpecType,
    load_override_fn: Optional[run_inference.LoadOverrideFnType] = None
) -> beam.pvalue.PCollection

This API is experimental and may change in the future.

Supports the same inference specs as RunInference. Inputs must consist of a keyed list of examples, and outputs consist of keyed list of prediction logs corresponding by index.

examples A PCollection of keyed, batched inputs of type Example, SequenceExample, or bytes. Each type support inference specs corresponding to the unbatched cases described in RunInference. Supports

  • PCollection[Tuple[K, List[Example]]]
  • PCollection[Tuple[K, List[SequenceExample]]]
  • PCollection[Tuple[K, List[Bytes]]]
inference_spec_type Model inference endpoint.
load_override_fn Optional function taking a model path and sequence of tags, and returning a tf SavedModel. The loaded model must be equivalent in interface to the model that would otherwise be loaded. It is up to the caller to ensure compatibility. This argument is experimental and subject to change.

A PCollection of Tuple[K, List[PredictionLog]].