apache_beam.ml.inference.base module

An extensible run inference transform.

Users of this module can extend the ModelHandler class for any machine learning framework. A ModelHandler implementation is a required parameter of RunInference.

The transform will handle standard inference functionality like metric collection, sharing model between threads and batching elements.

class apache_beam.ml.inference.base.PredictionResult(example, inference)

Bases: tuple

Create new instance of PredictionResult(example, inference)

example

Alias for field number 0

inference

Alias for field number 1

class apache_beam.ml.inference.base.ModelHandler[source]

Bases: typing.Generic

Has the ability to load and apply an ML model.

load_model() → ModelT[source]

Loads and initializes a model for processing.

run_inference(batch: Sequence[ExampleT], model: ModelT, inference_args: Optional[Dict[str, Any]] = None) → Iterable[PredictionT][source]

Runs inferences on a batch of examples.

Parameters:
  • batch – A sequence of examples or features.
  • model – The model used to make inferences.
Returns:

An Iterable of Predictions.
get_num_bytes(batch: Sequence[ExampleT]) → int[source]
Returns:The number of bytes of data for a batch.
get_metrics_namespace() → str[source]
Returns:A namespace for metrics collected by RunInference transform.
get_resource_hints() → dict[source]
Returns:Resource hints for the transform.
batch_elements_kwargs() → Mapping[str, Any][source]
Returns:kwargs suitable for beam.BatchElements.
class apache_beam.ml.inference.base.KeyedModelHandler(unkeyed: apache_beam.ml.inference.base.ModelHandler[~ExampleT, ~PredictionT, ~ModelT][ExampleT, PredictionT, ModelT])[source]

Bases: apache_beam.ml.inference.base.ModelHandler

A ModelHandler that takes keyed examples and returns keyed predictions.

For example, if the original model was used with RunInference to take a PCollection[E] to a PCollection[P], this would take a PCollection[Tuple[K, E]] to a PCollection[Tuple[K, P]], allowing one to associate the outputs with the inputs based on the key.

load_model() → ModelT[source]
run_inference(batch: Sequence[Tuple[KeyT, ExampleT]], model: ModelT, inference_args: Optional[Dict[str, Any]] = None) → Iterable[Tuple[KeyT, PredictionT]][source]
get_num_bytes(batch: Sequence[Tuple[KeyT, ExampleT]]) → int[source]
get_metrics_namespace() → str[source]
get_resource_hints()[source]
batch_elements_kwargs()[source]
class apache_beam.ml.inference.base.MaybeKeyedModelHandler(unkeyed: apache_beam.ml.inference.base.ModelHandler[~ExampleT, ~PredictionT, ~ModelT][ExampleT, PredictionT, ModelT])[source]

Bases: apache_beam.ml.inference.base.ModelHandler

A ModelHandler that takes possibly keyed examples and returns possibly keyed predictions.

For example, if the original model was used with RunInference to take a PCollection[E] to a PCollection[P], this would take either PCollection[E] to a PCollection[P] or PCollection[Tuple[K, E]] to a PCollection[Tuple[K, P]], depending on the whether the elements happen to be tuples, allowing one to associate the outputs with the inputs based on the key.

Note that this cannot be used if E happens to be a tuple type. In addition, either all examples should be keyed, or none of them.

load_model() → ModelT[source]
run_inference(batch: Sequence[Union[ExampleT, Tuple[KeyT, ExampleT]]], model: ModelT, inference_args: Optional[Dict[str, Any]] = None) → Union[Iterable[PredictionT], Iterable[Tuple[KeyT, PredictionT]]][source]
get_num_bytes(batch: Sequence[Union[ExampleT, Tuple[KeyT, ExampleT]]]) → int[source]
get_metrics_namespace() → str[source]
get_resource_hints()[source]
batch_elements_kwargs()[source]
class apache_beam.ml.inference.base.RunInference(model_handler: apache_beam.ml.inference.base.ModelHandler[~ExampleT, ~PredictionT, typing.Any][ExampleT, PredictionT, Any], clock=<module 'time' (built-in)>, inference_args: Optional[Dict[str, Any]] = None)[source]

Bases: apache_beam.transforms.ptransform.PTransform

A transform that takes a PCollection of examples (or features) to be used on an ML model. It will then output inferences (or predictions) for those examples in a PCollection of PredictionResults, containing the input examples and output inferences.

Models for supported frameworks can be loaded via a URI. Supported services can also be used.

Parameters:
  • model_handler – An implementation of ModelHandler.
  • clock – A clock implementing time_ns.
  • inference_args – Extra arguments for models whose inference call requires extra parameters.
classmethod create(model_handler_provider, **kwargs)[source]

Multi-language friendly constructor.

This constructor can be used with fully_qualified_named_transform to initialize RunInference transform from PythonCallableSource provided by foreign SDKs.

Parameters:
  • model_handler_provider – A callable object that returns ModelHandler.
  • kwargs – Keyword arguments for model_handler_provider.
expand(pcoll: apache_beam.pvalue.PCollection[~ExampleT][ExampleT]) → apache_beam.pvalue.PCollection[~PredictionT][PredictionT][source]