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

Bases: tuple

A NamedTuple containing both input and output from the inference.

Create new instance of PredictionResult(example, inference)

example

The input example.

inference

Results for the inference on the model for the given example.

class apache_beam.ml.inference.base.ModelHandler

Bases: typing.Generic

Has the ability to load and apply an ML model.

load_model() → ModelT

Loads and initializes a model for processing.

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

Runs inferences on a batch of examples.

Parameters:

• batch – A sequence of examples or features.
• model – The model used to make inferences.
• inference_args – Extra arguments for models whose inference call requires extra parameters.

Returns:

An Iterable of Predictions.

get_num_bytes(batch: Sequence[ExampleT]) → int

Returns:The number of bytes of data for a batch.

get_metrics_namespace() → str

Returns:A namespace for metrics collected by RunInference transform.

get_resource_hints() → dict

Returns:Resource hints for the transform.

batch_elements_kwargs() → Mapping[str, Any]

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])

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.

Parameters:unkeyed – An implementation of ModelHandler that does not require keys.

load_model() → ModelT

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

get_num_bytes(batch: Sequence[Tuple[KeyT, ExampleT]]) → int

get_metrics_namespace() → str

get_resource_hints()

batch_elements_kwargs()

class apache_beam.ml.inference.base.MaybeKeyedModelHandler(unkeyed: apache_beam.ml.inference.base.ModelHandler[~ExampleT, ~PredictionT, ~ModelT][ExampleT, PredictionT, ModelT])

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.

Parameters:unkeyed – An implementation of ModelHandler that does not require keys.

load_model() → ModelT

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]]]

get_num_bytes(batch: Sequence[Union[ExampleT, Tuple[KeyT, ExampleT]]]) → int

get_metrics_namespace() → str

get_resource_hints()

batch_elements_kwargs()

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)

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.

This transform attempts to batch examples using the beam.BatchElements transform. Batching may be configured using the ModelHandler.

Parameters:

• model_handler – An implementation of ModelHandler.
• clock – A clock implementing time_ns. Used for unit testing.
• inference_args – Extra arguments for models whose inference call requires extra parameters.

classmethod from_callable(model_handler_provider, **kwargs)

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]