#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import abc

import functools

import logging

import os

import tempfile

import uuid

from collections.abc import Callable

from collections.abc import Mapping

from collections.abc import Sequence

from dataclasses import dataclass

from typing import Any

from typing import Dict

from typing import Generic

from typing import Iterable

from typing import List

from typing import Optional

from typing import TypeVar

from typing import Union

from typing import cast

import jsonpickle

import numpy as np

import apache_beam as beam

from apache_beam.io.filesystems import FileSystems

from apache_beam.metrics.metric import Metrics

from apache_beam.ml.inference.base import ModelHandler

from apache_beam.ml.inference.base import ModelT

from apache_beam.ml.inference.base import RunInferenceDLQ

from apache_beam.options.pipeline_options import PipelineOptions

_LOGGER = logging.getLogger(__name__)

_ATTRIBUTE_FILE_NAME = 'attributes.json'

__all__ = [

'MLTransform',

'ProcessHandler',

'MLTransformProvider',

'BaseOperation',

'EmbeddingsManager'

]

TransformedDatasetT = TypeVar('TransformedDatasetT')

TransformedMetadataT = TypeVar('TransformedMetadataT')

# Input/Output types to the MLTransform.

MLTransformOutputT = TypeVar('MLTransformOutputT')

ExampleT = TypeVar('ExampleT')

# Input to the apply() method of BaseOperation.

OperationInputT = TypeVar('OperationInputT')

# Output of the apply() method of BaseOperation.

OperationOutputT = TypeVar('OperationOutputT')

# Input to the EmbeddingTypeAdapter input_fn

EmbeddingTypeAdapterInputT = TypeVar(

'EmbeddingTypeAdapterInputT') # e.g., Chunk

# Output of the EmbeddingTypeAdapter output_fn

EmbeddingTypeAdapterOutputT = TypeVar(

'EmbeddingTypeAdapterOutputT') # e.g., Embedding

@dataclass

class EmbeddingTypeAdapter(Generic[EmbeddingTypeAdapterInputT,

EmbeddingTypeAdapterOutputT]):

"""Adapts input types to text for embedding and converts output embeddings.

Args:

input_fn: Function to extract text for embedding from input type

output_fn: Function to create output type from input and embeddings

"""

input_fn: Callable[[Sequence[EmbeddingTypeAdapterInputT]], List[str]]

output_fn: Callable[[Sequence[EmbeddingTypeAdapterInputT], Sequence[Any]],

List[EmbeddingTypeAdapterOutputT]]

def __reduce__(self):

"""Custom serialization that preserves type information during

jsonpickle."""

return (self.__class__, (self.input_fn, self.output_fn))

def _map_errors_to_beam_row(element, cls_name=None):

row_elements = {

'element': element[0],

'msg': str(element[1][1]),

'stack': str(element[1][2]),

}

if cls_name is not None:

row_elements['transform_name'] = cls_name

return beam.Row(**row_elements)

class ArtifactMode(object):

PRODUCE = 'produce'

CONSUME = 'consume'

class MLTransformProvider:

"""

Data processing transforms that are intended to be used with MLTransform

should subclass MLTransformProvider and implement

get_ptransform_for_processing().

get_ptransform_for_processing() method should return a PTransform that can be

used to process the data.

"""

@abc.abstractmethod

def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform:

"""

Returns a PTransform that can be used to process the data.

"""

def get_counter(self):

"""

Returns the counter name for the data processing transform.

"""

counter_name = self.__class__.__name__

return Metrics.counter(MLTransform, f'BeamML_{counter_name}')

class BaseOperation(Generic[OperationInputT, OperationOutputT],

MLTransformProvider,

abc.ABC):

def __init__(self, columns: list[str]) -> None:

"""

Base Opertation class data processing transformations.

Args:

columns: List of column names to apply the transformation.

"""

self.columns = columns

@abc.abstractmethod

def apply_transform(self, data: OperationInputT,

output_column_name: str) -> dict[str, OperationOutputT]:

"""

Define any processing logic in the apply_transform() method.

processing logics are applied on inputs and returns a transformed

output.

Args:

inputs: input data.

"""

def __call__(self, data: OperationInputT,

output_column_name: str) -> dict[str, OperationOutputT]:

"""

This method is called when the instance of the class is called.

This method will invoke the apply() method of the class.

"""

transformed_data = self.apply_transform(data, output_column_name)

return transformed_data

class ProcessHandler(

beam.PTransform[beam.PCollection[ExampleT],

Union[beam.PCollection[MLTransformOutputT],

tuple[beam.PCollection[MLTransformOutputT],

beam.PCollection[beam.Row]]]],

abc.ABC):

"""

Only for internal use. No backwards compatibility guarantees.

"""

@abc.abstractmethod

def append_transform(self, transform: BaseOperation):

"""

Append transforms to the ProcessHandler.

"""

def _dict_input_fn(

columns: Sequence[str], batch: Sequence[Union[Dict[str, Any],

beam.Row]]) -> List[str]:

"""Extract text from specified columns in batch."""

if batch and hasattr(batch[0], '_asdict'):

batch = [row._asdict() if hasattr(row, '_asdict') else row for row in batch]

if not batch or not isinstance(batch[0], dict):

raise TypeError(

'Expected data to be dicts, got '

f'{type(batch[0])} instead.')

result = []

expected_keys = set(batch[0].keys())

expected_columns = set(columns)

# Process one batch item at a time

for item in batch:

item_keys = item.keys() if isinstance(item, dict) else set()

if set(item_keys) != expected_keys:

extra_keys = item_keys - expected_keys

missing_keys = expected_keys - item_keys

raise RuntimeError(

f'All dicts in batch must have the same keys. '

f'extra keys: {extra_keys}, '

f'missing keys: {missing_keys}')

missing_columns = expected_columns - item_keys

if (missing_columns):

raise RuntimeError(

f'Data does not contain the following columns '

f': {missing_columns}.')

# Get all columns for this item

for col in columns:

if isinstance(item, dict):

result.append(item[col])

return result

def _dict_output_fn(

columns: Sequence[str],

batch: Sequence[Union[Dict[str, Any], beam.Row]],

embeddings: Sequence[Any]) -> list[Union[dict[str, Any], beam.Row]]:

"""Map embeddings back to columns in batch."""

is_beam_row = False

if batch and hasattr(batch[0], '_asdict'):

is_beam_row = True

batch = [row._asdict() if hasattr(row, '_asdict') else row for row in batch]

result = []

for batch_idx, item in enumerate(batch):

for col_idx, col in enumerate(columns):

embedding_idx = batch_idx * len(columns) + col_idx

if isinstance(item, dict):

item[col] = embeddings[embedding_idx]

result.append(item)

if is_beam_row:

result = [beam.Row(**item) for item in result if isinstance(item, dict)]

return result

def _create_dict_adapter(

columns: List[str]) -> EmbeddingTypeAdapter[Dict[str, Any], Dict[str, Any]]:

"""Create adapter for dict-based processing."""

return EmbeddingTypeAdapter[Dict[str, Any], Dict[str, Any]](

input_fn=cast(

Callable[[Sequence[Dict[str, Any]]], List[str]],

functools.partial(_dict_input_fn, columns)),

output_fn=cast(

Callable[[Sequence[Dict[str, Any]], Sequence[Any]],

List[Dict[str, Any]]],

functools.partial(_dict_output_fn, columns)))

# TODO:https://github.com/apache/beam/issues/29356

# Add support for inference_fn

class EmbeddingsManager(MLTransformProvider):

def __init__(

self,

*,

columns: Optional[list[str]] = None,

type_adapter: Optional[EmbeddingTypeAdapter] = None,

# common args for all ModelHandlers.

load_model_args: Optional[dict[str, Any]] = None,

min_batch_size: Optional[int] = None,

max_batch_size: Optional[int] = None,

large_model: bool = False,

**kwargs):

self.load_model_args = load_model_args or {}

self.min_batch_size = min_batch_size

self.max_batch_size = max_batch_size

self.large_model = large_model

self.columns = columns

if columns is not None:

self.type_adapter = _create_dict_adapter(columns)

elif type_adapter is not None:

self.type_adapter = type_adapter

else:

raise ValueError("Either columns or type_adapter must be specified")

self.inference_args = kwargs.pop('inference_args', {})

if kwargs:

_LOGGER.warning("Ignoring the following arguments: %s", kwargs.keys())

# TODO:https://github.com/apache/beam/pull/29564 add set_model_handler method

@abc.abstractmethod

def get_model_handler(self) -> ModelHandler:

"""

Return framework specific model handler.

"""

def get_columns_to_apply(self):

return self.columns

class MLTransform(

beam.PTransform[beam.PCollection[ExampleT],

Union[beam.PCollection[MLTransformOutputT],

tuple[beam.PCollection[MLTransformOutputT],

beam.PCollection[beam.Row]]]],

Generic[ExampleT, MLTransformOutputT]):

def __init__(

self,

*,

write_artifact_location: Optional[str] = None,

read_artifact_location: Optional[str] = None,

transforms: Optional[list[MLTransformProvider]] = None):

"""

MLTransform is a Beam PTransform that can be used to apply

transformations to the data. MLTransform is used to wrap the

data processing transforms provided by Apache Beam. MLTransform

works in two modes: write and read. In the write mode,

MLTransform will apply the transforms to the data and store the

artifacts in the write_artifact_location. In the read mode,

MLTransform will read the artifacts from the

read_artifact_location and apply the transforms to the data. The

artifact location should be a valid storage path where the artifacts

can be written to or read from.

Note that when consuming artifacts, it is not necessary to pass the

transforms since they are inherently stored within the artifacts

themselves.

Args:

write_artifact_location: A storage location for artifacts resulting from

MLTransform. These artifacts include transformations applied to

the dataset and generated values like min, max from ScaleTo01,

and mean, var from ScaleToZScore. Artifacts are produced and written

to this location when using `write_artifact_mode`.

Later MLTransforms can reuse produced artifacts by setting

`read_artifact_mode` instead of `write_artifact_mode`. The value

assigned to `write_artifact_location` should be a valid storage

directory that the artifacts from this transform can be written to.

If no directory exists at this location, one will be created. This will

overwrite any artifacts already in this location, so distinct locations

should be used for each instance of MLTransform. Only one of

write_artifact_location and read_artifact_location should be specified.

read_artifact_location: A storage location to read artifacts resulting

froma previous MLTransform. These artifacts include transformations

applied to the dataset and generated values like min, max from

ScaleTo01, and mean, var from ScaleToZScore. Note that when consuming

artifacts, it is not necessary to pass the transforms since they are

inherently stored within the artifacts themselves. The value assigned

to `read_artifact_location` should be a valid storage path where the

artifacts can be read from. Only one of write_artifact_location and

read_artifact_location should be specified.

transforms: A list of transforms to apply to the data. All the transforms

are applied in the order they are specified. The input of the

i-th transform is the output of the (i-1)-th transform. Multi-input

transforms are not supported yet.

"""

if read_artifact_location and write_artifact_location:

raise ValueError(

'Only one of read_artifact_location or write_artifact_location can '

'be specified to initialize MLTransform')

if not read_artifact_location and not write_artifact_location:

raise ValueError(

'Either a read_artifact_location or write_artifact_location must be '

'specified to initialize MLTransform')

if read_artifact_location:

artifact_location = read_artifact_location

artifact_mode = ArtifactMode.CONSUME

if transforms:

raise ValueError(

'Transforms should not be passed in read mode. In read mode, '

'the transforms are read from the artifact location.')

else:

artifact_location = write_artifact_location # type: ignore[assignment]

artifact_mode = ArtifactMode.PRODUCE

self._parent_artifact_location = artifact_location

self._artifact_mode = artifact_mode

self.transforms = transforms or []

self._counter = Metrics.counter(

MLTransform, f'BeamML_{self.__class__.__name__}')

self._with_exception_handling = False

self._exception_handling_args: dict[str, Any] = {}

def expand(

self, pcoll: beam.PCollection[ExampleT]

) -> Union[beam.PCollection[MLTransformOutputT],

tuple[beam.PCollection[MLTransformOutputT],

beam.PCollection[beam.Row]]]:

"""

This is the entrypoint for the MLTransform. This method will

invoke the process_data() method of the ProcessHandler instance

to process the incoming data.

process_data takes in a PCollection and applies the PTransforms

necessary to process the data and returns a PCollection of

transformed data.

Args:

pcoll: A PCollection of ExampleT type.

Returns:

A PCollection of MLTransformOutputT type

"""

upstream_errors = []

_ = [self._validate_transform(transform) for transform in self.transforms]

if self._artifact_mode == ArtifactMode.PRODUCE:

ptransform_partitioner = _MLTransformToPTransformMapper(

transforms=self.transforms,

artifact_location=self._parent_artifact_location,

artifact_mode=self._artifact_mode,

pipeline_options=pcoll.pipeline.options)

ptransform_list = ptransform_partitioner.create_and_save_ptransform_list()

else:

ptransform_list = (

_MLTransformToPTransformMapper.load_transforms_from_artifact_location(

self._parent_artifact_location))

# the saved transforms has artifact mode set to PRODUCE.

# set the artifact mode to CONSUME.

for i in range(len(ptransform_list)):

if hasattr(ptransform_list[i], 'artifact_mode'):

ptransform_list[i].artifact_mode = self._artifact_mode

transform_name = None

for ptransform in ptransform_list:

if self._with_exception_handling:

if hasattr(ptransform, 'with_exception_handling'):

ptransform = ptransform.with_exception_handling(

**self._exception_handling_args)

pcoll, bad_results = pcoll | ptransform

# RunInference outputs a RunInferenceDLQ instead of a PCollection.

# since TFTProcessHandler and RunInferene are supported, try to infer

# the type of bad_results and append it to the list of errors.

if isinstance(bad_results, RunInferenceDLQ):

bad_results = bad_results.failed_inferences

transform_name = ptransform.annotations()['model_handler']

elif not isinstance(bad_results, beam.PCollection):

raise NotImplementedError(

f'Unexpected type for bad_results: {type(bad_results)}')

bad_results = bad_results | beam.Map(

lambda x: _map_errors_to_beam_row(x, transform_name))

upstream_errors.append(bad_results)

else:

pcoll = pcoll | ptransform

_ = (

pcoll.pipeline

| "MLTransformMetricsUsage" >> MLTransformMetricsUsage(self))

if self._with_exception_handling:

bad_pcoll = (upstream_errors | beam.Flatten())

return pcoll, bad_pcoll # type: ignore[return-value]

return pcoll # type: ignore[return-value]

def with_transform(self, transform: MLTransformProvider):

"""

Add a transform to the MLTransform pipeline.

Args:

transform: A BaseOperation instance.

Returns:

A MLTransform instance.

"""

self._validate_transform(transform)

self.transforms.append(transform)

return self

def _validate_transform(self, transform):

if not isinstance(transform, MLTransformProvider):

raise TypeError(

'transform must be a subclass of BaseOperation. '

'Got: %s instead.' % type(transform))

def with_exception_handling(

self, *, exc_class=Exception, use_subprocess=False, threshold=1):

self._with_exception_handling = True

self._exception_handling_args = {

'exc_class': exc_class,

'use_subprocess': use_subprocess,

'threshold': threshold

}

return self

class MLTransformMetricsUsage(beam.PTransform):

def __init__(self, ml_transform: MLTransform):

self._ml_transform = ml_transform

self._ml_transform._counter.inc()

def expand(self, pipeline):

def _increment_counters():

# increment for MLTransform.

self._ml_transform._counter.inc()

# increment if data processing transforms are passed.

transforms = self._ml_transform.transforms

if transforms:

for transform in transforms:

transform.get_counter().inc()

_ = (

pipeline

| beam.Create([None])

| beam.Map(lambda _: _increment_counters()))

class _TransformAttributeManager:

"""

Base class used for saving and loading the attributes.

"""

@staticmethod

def save_attributes(artifact_location):

"""

Save the attributes to json file using stdlib json.

"""

raise NotImplementedError

@staticmethod

def load_attributes(artifact_location):

"""

Load the attributes from json file.

"""

raise NotImplementedError

class _JsonPickleTransformAttributeManager(_TransformAttributeManager):

"""

Use Jsonpickle to save and load the attributes. Here the attributes refer

to the list of PTransforms that are used to process the data.

jsonpickle is used to serialize the PTransforms and save it to a json file and

is compatible across python versions.

"""

@staticmethod

def _is_remote_path(path):

is_gcs = path.find('gs://') != -1

# TODO:https://github.com/apache/beam/issues/29356

# Add support for other remote paths.

if not is_gcs and path.find('://') != -1:

raise RuntimeError(

"Artifact locations are currently supported for only available for "

"local paths and GCS paths. Got: %s" % path)

return is_gcs

@staticmethod

def save_attributes(

ptransform_list,

artifact_location,

**kwargs,

):

# if an artifact location is present, instead of overwriting the

# existing file, raise an error since the same artifact location

# can be used by multiple beam jobs and this could result in undesired

# behavior.

if FileSystems.exists(FileSystems.join(artifact_location,

_ATTRIBUTE_FILE_NAME)):

raise FileExistsError(

"The artifact location %s already exists and contains %s. Please "

"specify a different location." %

(artifact_location, _ATTRIBUTE_FILE_NAME))

if _JsonPickleTransformAttributeManager._is_remote_path(artifact_location):

temp_dir = tempfile.mkdtemp()

temp_json_file = os.path.join(temp_dir, _ATTRIBUTE_FILE_NAME)

with open(temp_json_file, 'w+') as f:

f.write(jsonpickle.encode(ptransform_list))

with open(temp_json_file, 'rb') as f:

from apache_beam.runners.dataflow.internal import apiclient

_LOGGER.info('Creating artifact location: %s', artifact_location)

# pipeline options required to for the client to configure project.

options = kwargs.get('options')

try:

apiclient.DataflowApplicationClient(options=options).stage_file(

gcs_or_local_path=artifact_location,

file_name=_ATTRIBUTE_FILE_NAME,

stream=f,

mime_type='application/json')

except Exception as exc:

if not options:

raise RuntimeError(

"Failed to create Dataflow client. "

"Pipeline options are required to save the attributes."

"in the artifact location %s" % artifact_location) from exc

raise

else:

if not FileSystems.exists(artifact_location):

FileSystems.mkdirs(artifact_location)

# FileSystems.open() fails if the file does not exist.

with open(os.path.join(artifact_location, _ATTRIBUTE_FILE_NAME),

'w+') as f:

f.write(jsonpickle.encode(ptransform_list))

@staticmethod

def load_attributes(artifact_location):

with FileSystems.open(os.path.join(artifact_location, _ATTRIBUTE_FILE_NAME),

'rb') as f:

return jsonpickle.decode(f.read())

_transform_attribute_manager = _JsonPickleTransformAttributeManager

class _MLTransformToPTransformMapper:

"""

This class takes in a list of data processing transforms compatible to be

wrapped around MLTransform and returns a list of PTransforms that are used to

run the data processing transforms.

The _MLTransformToPTransformMapper is responsible for loading and saving the

PTransforms or attributes of PTransforms to the artifact location to seal

the gap between the training and inference pipelines.

"""

def __init__(

self,

transforms: list[MLTransformProvider],

artifact_location: str,

artifact_mode: str = ArtifactMode.PRODUCE,

pipeline_options: Optional[PipelineOptions] = None,

):

self.transforms = transforms

self._parent_artifact_location = artifact_location

self.artifact_mode = artifact_mode

self.pipeline_options = pipeline_options

def create_and_save_ptransform_list(self):

ptransform_list = self.create_ptransform_list()

self.save_transforms_in_artifact_location(ptransform_list)

return ptransform_list

def create_ptransform_list(self):

previous_ptransform_type = None

current_ptransform = None

ptransform_list = []

for transform in self.transforms:

if not isinstance(transform, MLTransformProvider):

raise RuntimeError(

'Transforms must be instances of MLTransformProvider and '

'implement get_ptransform_for_processing() method.')

# for each instance of PTransform, create a new artifact location

current_ptransform = transform.get_ptransform_for_processing(

artifact_location=os.path.join(

self._parent_artifact_location, uuid.uuid4().hex[:6]),

artifact_mode=self.artifact_mode)

append_transform = hasattr(current_ptransform, 'append_transform')

if (type(current_ptransform)

!= previous_ptransform_type) or not append_transform:

ptransform_list.append(current_ptransform)

previous_ptransform_type = type(current_ptransform)

# If different PTransform is appended to the list and the PTransform

# supports append_transform, append the transform to the PTransform.

if append_transform:

ptransform_list[-1].append_transform(transform)

return ptransform_list

def save_transforms_in_artifact_location(self, ptransform_list):

"""

Save the ptransform references to json file.

"""

_transform_attribute_manager.save_attributes(

ptransform_list=ptransform_list,

artifact_location=self._parent_artifact_location,

options=self.pipeline_options)

@staticmethod

def load_transforms_from_artifact_location(artifact_location):

return _transform_attribute_manager.load_attributes(artifact_location)

class _EmbeddingHandler(ModelHandler):

"""

A ModelHandler intended to be work on list[dict[str, Any]] inputs.

The inputs to the model handler are expected to be a list of dicts.

For example, if the original mode is used with RunInference to take a

PCollection[E] to a PCollection[P], this ModelHandler would take a

PCollection[dict[str, E]] to a PCollection[dict[str, P]].

_EmbeddingHandler will accept an EmbeddingsManager instance, which

contains the details of the model to be loaded and the inference_fn to be

used. The purpose of _EmbeddingHandler is to generate embeddings for

general inputs using the EmbeddingsManager instance.

This is an internal class and offers no backwards compatibility guarantees.

Args:

embeddings_manager: An EmbeddingsManager instance.

"""

def __init__(self, embeddings_manager: EmbeddingsManager):

self.embedding_config = embeddings_manager

self._underlying = self.embedding_config.get_model_handler()

self.columns = self.embedding_config.get_columns_to_apply()

def load_model(self):

model = self._underlying.load_model()

return model

def _validate_column_data(self, batch):

pass

def run_inference(

self,

batch: Sequence[dict[str, list[str]]],

model: ModelT,

inference_args: Optional[dict[str, Any]] = None,

) -> list[dict[str, Union[list[float], list[str]]]]:

"""

Runs inference on a batch of text inputs. The inputs are expected to be

a list of dicts. Each dict should have the same keys, and the shape

should be of the same size for a single key across the batch.

"""

embedding_input = self.embedding_config.type_adapter.input_fn(batch)

self._validate_column_data(batch=embedding_input)

prediction = self._underlying.run_inference(

embedding_input, model, inference_args)

# Convert prediction to Sequence[Any]

if isinstance(prediction, np.ndarray):

prediction_seq = prediction.tolist()

elif isinstance(prediction, Iterable) and not isinstance(prediction,

(str, bytes)):

prediction_seq = list(prediction)

else:

prediction_seq = [prediction]

return self.embedding_config.type_adapter.output_fn(batch, prediction_seq)

def get_metrics_namespace(self) -> str:

return (

self._underlying.get_metrics_namespace() or 'BeamML_EmbeddingHandler')

def batch_elements_kwargs(self) -> Mapping[str, Any]:

batch_sizes_map = {}

if self.embedding_config.max_batch_size:

batch_sizes_map['max_batch_size'] = self.embedding_config.max_batch_size

if self.embedding_config.min_batch_size:

batch_sizes_map['min_batch_size'] = self.embedding_config.min_batch_size

return (self._underlying.batch_elements_kwargs() or batch_sizes_map)

def __repr__(self):

return self._underlying.__repr__()

def validate_inference_args(self, _):

pass

class _TextEmbeddingHandler(_EmbeddingHandler):

"""

A ModelHandler intended to be work on list[dict[str, str]] inputs.

The inputs to the model handler are expected to be a list of dicts.

For example, if the original mode is used with RunInference to take a

PCollection[E] to a PCollection[P], this ModelHandler would take a

PCollection[dict[str, E]] to a PCollection[dict[str, P]].

_TextEmbeddingHandler will accept an EmbeddingsManager instance, which

contains the details of the model to be loaded and the inference_fn to be

used. The purpose of _TextEmbeddingHandler is to generate embeddings for

text inputs using the EmbeddingsManager instance.

If the input is not a text column, a RuntimeError will be raised.

This is an internal class and offers no backwards compatibility guarantees.

Args:

embeddings_manager: An EmbeddingsManager instance.

"""

def _validate_column_data(self, batch):

if not isinstance(batch[0], (str, bytes)):

raise TypeError(

'Embeddings can only be generated on dict[str, str].'

f'Got dict[str, {type(batch[0])}] instead.')

def get_metrics_namespace(self) -> str:

return (

self._underlying.get_metrics_namespace() or

'BeamML_TextEmbeddingHandler')

class _ImageEmbeddingHandler(_EmbeddingHandler):

"""

A ModelHandler intended to be work on list[dict[str, Image]] inputs.

The inputs to the model handler are expected to be a list of dicts.

For example, if the original mode is used with RunInference to take a

PCollection[E] to a PCollection[P], this ModelHandler would take a

PCollection[dict[str, E]] to a PCollection[dict[str, P]].

_ImageEmbeddingHandler will accept an EmbeddingsManager instance, which

contains the details of the model to be loaded and the inference_fn to be

used. The purpose of _ImageEmbeddingHandler is to generate embeddings for

image inputs using the EmbeddingsManager instance.

If the input is not an Image representation column, a RuntimeError will be

raised.

This is an internal class and offers no backwards compatibility guarantees.

Args:

embeddings_manager: An EmbeddingsManager instance.

"""

def _validate_column_data(self, batch):

# Don't want to require framework-specific imports

# here, so just catch columns of primatives for now.

if isinstance(batch[0], (int, str, float, bool)):

raise TypeError(

'Embeddings can only be generated on dict[str, Image].'

f'Got dict[str, {type(batch[0])}] instead.')

def get_metrics_namespace(self) -> str:

return (

self._underlying.get_metrics_namespace() or

'BeamML_ImageEmbeddingHandler')

class _MultiModalEmbeddingHandler(_EmbeddingHandler):

"""

A ModelHandler intended to be work on

list[dict[str, TypedDict(Image, Video, str)]] inputs.

The inputs to the model handler are expected to be a list of dicts.

For example, if the original mode is used with RunInference to take a

PCollection[E] to a PCollection[P], this ModelHandler would take a

PCollection[dict[str, E]] to a PCollection[dict[str, P]].

_MultiModalEmbeddingHandler will accept an EmbeddingsManager instance, which

contains the details of the model to be loaded and the inference_fn to be

used. The purpose of _MultiMOdalEmbeddingHandler is to generate embeddings

for image, video, and text inputs using the EmbeddingsManager instance.

If the input is not an Image representation column, a RuntimeError will be

raised.

This is an internal class and offers no backwards compatibility guarantees.

Args:

embeddings_manager: An EmbeddingsManager instance.

"""

def _validate_column_data(self, batch):

# Don't want to require framework-specific imports

# here, so just catch columns of primatives for now.

if isinstance(batch[0], (int, str, float, bool)):

raise TypeError(

'Embeddings can only be generated on '

' dict[str, dataclass] types. '

f'Got dict[str, {type(batch[0])}] instead.')

def get_metrics_namespace(self) -> str:

return (

self._underlying.get_metrics_namespace() or

'BeamML_MultiModalEmbeddingHandler')