| # |
| # 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') |
