#
# 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.
#
# pytype: skip-file
import logging
import sys
from collections import defaultdict
from enum import Enum
from typing import Any
from typing import Callable
from typing import Dict
from typing import Iterable
from typing import Optional
from typing import Sequence
from typing import Union
import tensorflow as tf
import torch
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import ModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.pytorch_inference import _convert_to_device
from transformers import AutoModel
from transformers import Pipeline
from transformers import TFAutoModel
from transformers import pipeline
_LOGGER = logging.getLogger(__name__)
__all__ = [
"HuggingFaceModelHandlerTensor",
"HuggingFaceModelHandlerKeyedTensor",
"HuggingFacePipelineModelHandler",
]
TensorInferenceFn = Callable[[
Sequence[Union[torch.Tensor, tf.Tensor]],
Union[AutoModel, TFAutoModel],
str,
Optional[Dict[str, Any]],
Optional[str],
],
Iterable[PredictionResult],
]
KeyedTensorInferenceFn = Callable[[
Sequence[Dict[str, Union[torch.Tensor, tf.Tensor]]],
Union[AutoModel, TFAutoModel],
str,
Optional[Dict[str, Any]],
Optional[str]
],
Iterable[PredictionResult]]
PipelineInferenceFn = Callable[
[Sequence[str], Pipeline, Optional[Dict[str, Any]]],
Iterable[PredictionResult]]
class PipelineTask(str, Enum):
"""
PipelineTask defines all the tasks supported by the Hugging Face Pipelines
listed at https://huggingface.co/docs/transformers/main_classes/pipelines.
Only these tasks can be passed to HuggingFacePipelineModelHandler.
"""
AudioClassification = 'audio-classification'
AutomaticSpeechRecognition = 'automatic-speech-recognition'
Conversational = 'conversational'
DepthEstimation = 'depth-estimation'
DocumentQuestionAnswering = 'document-question-answering'
FeatureExtraction = 'feature-extraction'
FillMask = 'fill-mask'
ImageClassification = 'image-classification'
ImageSegmentation = 'image-segmentation'
ImageToText = 'image-to-text'
MaskGeneration = 'mask-generation'
NER = 'ner'
ObjectDetection = 'object-detection'
QuestionAnswering = 'question-answering'
SentimentAnalysis = 'sentiment-analysis'
Summarization = 'summarization'
TableQuestionAnswering = 'table-question-answering'
TextClassification = 'text-classification'
TextGeneration = 'text-generation'
Text2TextGeneration = 'text2text-generation'
TextToAudio = 'text-to-audio'
TokenClassification = 'token-classification'
Translation = 'translation'
VideoClassification = 'video-classification'
VisualQuestionAnswering = 'visual-question-answering'
VQA = 'vqa'
ZeroShotAudioClassification = 'zero-shot-audio-classification'
ZeroShotClassification = 'zero-shot-classification'
ZeroShotImageClassification = 'zero-shot-image-classification'
ZeroShotObjectDetection = 'zero-shot-object-detection'
Translation_XX_to_YY = 'translation_XX_to_YY'
def _validate_constructor_args(model_uri, model_class):
message = (
"Please provide both model class and model uri to load the model."
"Got params as model_uri={model_uri} and "
"model_class={model_class}.")
if not model_uri and not model_class:
raise RuntimeError(
message.format(model_uri=model_uri, model_class=model_class))
elif not model_uri:
raise RuntimeError(
message.format(model_uri=model_uri, model_class=model_class))
elif not model_class:
raise RuntimeError(
message.format(model_uri=model_uri, model_class=model_class))
def no_gpu_available_warning():
_LOGGER.warning(
"HuggingFaceModelHandler specified a 'GPU' device, "
"but GPUs are not available. Switching to CPU.")
def is_gpu_available_torch():
if torch.cuda.is_available():
return True
else:
no_gpu_available_warning()
return False
def get_device_torch(device):
if device == "GPU" and is_gpu_available_torch():
return torch.device("cuda")
return torch.device("cpu")
def is_gpu_available_tensorflow(device):
gpu_devices = tf.config.list_physical_devices(device)
if len(gpu_devices) == 0:
no_gpu_available_warning()
return False
return True
def _validate_constructor_args_hf_pipeline(task, model):
if not task and not model:
raise RuntimeError(
'Please provide either task or model to the '
'HuggingFacePipelineModelHandler. If the model already defines the '
'task, no need to specify the task.')
def _run_inference_torch_keyed_tensor(
batch: Sequence[Dict[str, torch.Tensor]],
model: AutoModel,
device,
inference_args: Dict[str, Any],
model_id: Optional[str] = None) -> Iterable[PredictionResult]:
device = get_device_torch(device)
key_to_tensor_list = defaultdict(list)
# torch.no_grad() mitigates GPU memory issues
# https://github.com/apache/beam/issues/22811
with torch.no_grad():
for example in batch:
for key, tensor in example.items():
key_to_tensor_list[key].append(tensor)
key_to_batched_tensors = {}
for key in key_to_tensor_list:
batched_tensors = torch.stack(key_to_tensor_list[key])
batched_tensors = _convert_to_device(batched_tensors, device)
key_to_batched_tensors[key] = batched_tensors
predictions = model(**key_to_batched_tensors, **inference_args)
return utils._convert_to_result(batch, predictions, model_id)
def _run_inference_tensorflow_keyed_tensor(
batch: Sequence[Dict[str, tf.Tensor]],
model: TFAutoModel,
device,
inference_args: Dict[str, Any],
model_id: Optional[str] = None) -> Iterable[PredictionResult]:
if device == "GPU":
is_gpu_available_tensorflow(device)
key_to_tensor_list = defaultdict(list)
for example in batch:
for key, tensor in example.items():
key_to_tensor_list[key].append(tensor)
key_to_batched_tensors = {}
for key in key_to_tensor_list:
batched_tensors = tf.stack(key_to_tensor_list[key], axis=0)
key_to_batched_tensors[key] = batched_tensors
predictions = model(**key_to_batched_tensors, **inference_args)
return utils._convert_to_result(batch, predictions, model_id)
[docs]class HuggingFaceModelHandlerKeyedTensor(ModelHandler[Dict[str,
Union[tf.Tensor,
torch.Tensor]],
PredictionResult,
Union[AutoModel,
TFAutoModel]]):
def __init__(
self,
model_uri: str,
model_class: Union[AutoModel, TFAutoModel],
framework: str,
device: str = "CPU",
*,
inference_fn: Optional[Callable[..., Iterable[PredictionResult]]] = None,
load_model_args: Optional[Dict[str, Any]] = None,
inference_args: Optional[Dict[str, Any]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
large_model: bool = False,
**kwargs):
"""
Implementation of the ModelHandler interface for HuggingFace with
Keyed Tensors for PyTorch/Tensorflow backend.
Example Usage model::
pcoll | RunInference(HuggingFaceModelHandlerKeyedTensor(
model_uri="bert-base-uncased", model_class=AutoModelForMaskedLM,
framework='pt'))
Args:
model_uri (str): path to the pretrained model on the hugging face
models hub.
model_class: model class to load the repository from model_uri.
framework (str): Framework to use for the model. 'tf' for TensorFlow and
'pt' for PyTorch.
device: For torch tensors, specify device on which you wish to
run the model. Defaults to CPU.
inference_fn: the inference function to use during RunInference.
Default is _run_inference_torch_keyed_tensor or
_run_inference_tensorflow_keyed_tensor depending on the input type.
load_model_args (Dict[str, Any]): (Optional) Keyword arguments to provide
load options while loading models from Hugging Face Hub.
Defaults to None.
inference_args (Dict[str, Any]): (Optional) Non-batchable arguments
required as inputs to the model's inference function. Unlike Tensors
in `batch`, these parameters will not be dynamically batched.
Defaults to None.
min_batch_size: the minimum batch size to use when batching inputs.
max_batch_size: the maximum batch size to use when batching inputs.
large_model: set to true if your model is large enough to run into
memory pressure if you load multiple copies. Given a model that
consumes N memory and a machine with W cores and M memory, you should
set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
**Supported Versions:** HuggingFaceModelHandler supports
transformers>=4.18.0.
"""
self._model_uri = model_uri
self._model_class = model_class
self._device = device
self._inference_fn = inference_fn
self._model_config_args = load_model_args if load_model_args else {}
self._inference_args = inference_args if inference_args else {}
self._batching_kwargs = {}
self._env_vars = kwargs.get("env_vars", {})
if min_batch_size is not None:
self._batching_kwargs["min_batch_size"] = min_batch_size
if max_batch_size is not None:
self._batching_kwargs["max_batch_size"] = max_batch_size
self._large_model = large_model
self._framework = framework
_validate_constructor_args(
model_uri=self._model_uri, model_class=self._model_class)
[docs] def load_model(self):
"""Loads and initializes the model for processing."""
model = self._model_class.from_pretrained(
self._model_uri, **self._model_config_args)
if self._framework == 'pt':
if self._device == "GPU" and is_gpu_available_torch:
model.to(torch.device("cuda"))
if callable(getattr(model, 'requires_grad_', None)):
model.requires_grad_(False)
return model
[docs] def run_inference(
self,
batch: Sequence[Dict[str, Union[tf.Tensor, torch.Tensor]]],
model: Union[AutoModel, TFAutoModel],
inference_args: Optional[Dict[str, Any]] = None
) -> Iterable[PredictionResult]:
"""
Runs inferences on a batch of Keyed Tensors and returns an Iterable of
Tensors Predictions.
This method stacks the list of Tensors in a vectorized format to optimize
the inference call.
Args:
batch: A sequence of Keyed Tensors. These Tensors should be batchable,
as this method will call `tf.stack()`/`torch.stack()` and pass in
batched Tensors with dimensions (batch_size, n_features, etc.) into
the model's predict() function.
model: A Tensorflow/PyTorch model.
inference_args: Non-batchable arguments required as inputs to the
model's inference function. Unlike Tensors in `batch`,
these parameters will not be dynamically batched.
Returns:
An Iterable of type PredictionResult.
"""
inference_args = {} if not inference_args else inference_args
if self._inference_fn:
return self._inference_fn(
batch, model, self._device, inference_args, self._model_uri)
if self._framework == "tf":
return _run_inference_tensorflow_keyed_tensor(
batch, model, self._device, inference_args, self._model_uri)
else:
return _run_inference_torch_keyed_tensor(
batch, model, self._device, inference_args, self._model_uri)
[docs] def update_model_path(self, model_path: Optional[str] = None):
self._model_uri = model_path if model_path else self._model_uri
[docs] def get_num_bytes(
self, batch: Sequence[Union[tf.Tensor, torch.Tensor]]) -> int:
"""
Returns:
The number of bytes of data for the Tensors batch.
"""
if self._framework == "tf":
return sum(sys.getsizeof(element) for element in batch)
else:
return sum(
(el.element_size() for tensor in batch for el in tensor.values()))
[docs] def batch_elements_kwargs(self):
return self._batching_kwargs
[docs] def share_model_across_processes(self) -> bool:
return self._large_model
[docs] def get_metrics_namespace(self) -> str:
"""
Returns:
A namespace for metrics collected by the RunInference transform.
"""
return "BeamML_HuggingFaceModelHandler_KeyedTensor"
def _default_inference_fn_torch(
batch: Sequence[Union[tf.Tensor, torch.Tensor]],
model: Union[AutoModel, TFAutoModel],
device,
inference_args: Dict[str, Any],
model_id: Optional[str] = None) -> Iterable[PredictionResult]:
device = get_device_torch(device)
# torch.no_grad() mitigates GPU memory issues
# https://github.com/apache/beam/issues/22811
with torch.no_grad():
batched_tensors = torch.stack(batch)
batched_tensors = _convert_to_device(batched_tensors, device)
predictions = model(batched_tensors, **inference_args)
return utils._convert_to_result(batch, predictions, model_id)
def _default_inference_fn_tensorflow(
batch: Sequence[Union[tf.Tensor, torch.Tensor]],
model: Union[AutoModel, TFAutoModel],
device,
inference_args: Dict[str, Any],
model_id: Optional[str] = None) -> Iterable[PredictionResult]:
if device == "GPU":
is_gpu_available_tensorflow(device)
batched_tensors = tf.stack(batch, axis=0)
predictions = model(batched_tensors, **inference_args)
return utils._convert_to_result(batch, predictions, model_id)
[docs]class HuggingFaceModelHandlerTensor(ModelHandler[Union[tf.Tensor, torch.Tensor],
PredictionResult,
Union[AutoModel,
TFAutoModel]]):
def __init__(
self,
model_uri: str,
model_class: Union[AutoModel, TFAutoModel],
device: str = "CPU",
*,
inference_fn: Optional[Callable[..., Iterable[PredictionResult]]] = None,
load_model_args: Optional[Dict[str, Any]] = None,
inference_args: Optional[Dict[str, Any]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
large_model: bool = False,
**kwargs):
"""
Implementation of the ModelHandler interface for HuggingFace with
Tensors for PyTorch/Tensorflow backend.
Depending on the type of tensors, the model framework is determined
automatically.
Example Usage model:
pcoll | RunInference(HuggingFaceModelHandlerTensor(
model_uri="bert-base-uncased", model_class=AutoModelForMaskedLM))
Args:
model_uri (str): path to the pretrained model on the hugging face
models hub.
model_class: model class to load the repository from model_uri.
device: For torch tensors, specify device on which you wish to
run the model. Defaults to CPU.
inference_fn: the inference function to use during RunInference.
Default is _run_inference_torch_keyed_tensor or
_run_inference_tensorflow_keyed_tensor depending on the input type.
load_model_args (Dict[str, Any]): (Optional) keyword arguments to provide
load options while loading models from Hugging Face Hub.
Defaults to None.
inference_args (Dict[str, Any]): (Optional) Non-batchable arguments
required as inputs to the model's inference function. Unlike Tensors
in `batch`, these parameters will not be dynamically batched.
Defaults to None.
min_batch_size: the minimum batch size to use when batching inputs.
max_batch_size: the maximum batch size to use when batching inputs.
large_model: set to true if your model is large enough to run into
memory pressure if you load multiple copies. Given a model that
consumes N memory and a machine with W cores and M memory, you should
set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
**Supported Versions:** HuggingFaceModelHandler supports
transformers>=4.18.0.
"""
self._model_uri = model_uri
self._model_class = model_class
self._device = device
self._inference_fn = inference_fn
self._model_config_args = load_model_args if load_model_args else {}
self._inference_args = inference_args if inference_args else {}
self._batching_kwargs = {}
self._env_vars = kwargs.get("env_vars", {})
if min_batch_size is not None:
self._batching_kwargs["min_batch_size"] = min_batch_size
if max_batch_size is not None:
self._batching_kwargs["max_batch_size"] = max_batch_size
self._large_model = large_model
self._framework = ""
_validate_constructor_args(
model_uri=self._model_uri, model_class=self._model_class)
[docs] def load_model(self):
"""Loads and initializes the model for processing."""
model = self._model_class.from_pretrained(
self._model_uri, **self._model_config_args)
if callable(getattr(model, 'requires_grad_', None)):
model.requires_grad_(False)
return model
[docs] def run_inference(
self,
batch: Sequence[Union[tf.Tensor, torch.Tensor]],
model: Union[AutoModel, TFAutoModel],
inference_args: Optional[Dict[str, Any]] = None
) -> Iterable[PredictionResult]:
"""
Runs inferences on a batch of Tensors and returns an Iterable of
Tensors Predictions.
This method stacks the list of Tensors in a vectorized format to optimize
the inference call.
Args:
batch: A sequence of Tensors. These Tensors should be batchable, as
this method will call `tf.stack()`/`torch.stack()` and pass in
batched Tensors with dimensions (batch_size, n_features, etc.)
into the model's predict() function.
model: A Tensorflow/PyTorch model.
inference_args (Dict[str, Any]): Non-batchable arguments required as
inputs to the model's inference function. Unlike Tensors in `batch`,
these parameters will not be dynamically batched.
Returns:
An Iterable of type PredictionResult.
"""
inference_args = {} if not inference_args else inference_args
if not self._framework:
if isinstance(batch[0], tf.Tensor):
self._framework = "tf"
else:
self._framework = "pt"
if (self._framework == "pt" and self._device == "GPU" and
is_gpu_available_torch()):
model.to(torch.device("cuda"))
if self._inference_fn:
return self._inference_fn(
batch, model, inference_args, inference_args, self._model_uri)
if self._framework == "tf":
return _default_inference_fn_tensorflow(
batch, model, self._device, inference_args, self._model_uri)
else:
return _default_inference_fn_torch(
batch, model, self._device, inference_args, self._model_uri)
[docs] def update_model_path(self, model_path: Optional[str] = None):
self._model_uri = model_path if model_path else self._model_uri
[docs] def get_num_bytes(
self, batch: Sequence[Union[tf.Tensor, torch.Tensor]]) -> int:
"""
Returns:
The number of bytes of data for the Tensors batch.
"""
if self._framework == "tf":
return sum(sys.getsizeof(element) for element in batch)
else:
return sum(
(el.element_size() for tensor in batch for el in tensor.values()))
[docs] def batch_elements_kwargs(self):
return self._batching_kwargs
[docs] def share_model_across_processes(self) -> bool:
return self._large_model
[docs] def get_metrics_namespace(self) -> str:
"""
Returns:
A namespace for metrics collected by the RunInference transform.
"""
return 'BeamML_HuggingFaceModelHandler_Tensor'
def _convert_to_result(
batch: Iterable,
predictions: Union[Iterable, Dict[Any, Iterable]],
model_id: Optional[str] = None,
) -> Iterable[PredictionResult]:
return [
PredictionResult(x, y, model_id) for x, y in zip(batch, [predictions])
]
def _default_pipeline_inference_fn(
batch, pipeline, inference_args) -> Iterable[PredictionResult]:
predicitons = pipeline(batch, **inference_args)
return predicitons
[docs]class HuggingFacePipelineModelHandler(ModelHandler[str,
PredictionResult,
Pipeline]):
def __init__(
self,
task: Union[str, PipelineTask] = "",
model: str = "",
*,
inference_fn: PipelineInferenceFn = _default_pipeline_inference_fn,
load_pipeline_args: Optional[Dict[str, Any]] = None,
inference_args: Optional[Dict[str, Any]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
large_model: bool = False,
**kwargs):
"""
Implementation of the ModelHandler interface for Hugging Face Pipelines.
**Note:** To specify which device to use (CPU/GPU),
use the load_pipeline_args with key-value as you would do in the usual
Hugging Face pipeline. Ex: load_pipeline_args={'device':0})
Example Usage model::
pcoll | RunInference(HuggingFacePipelineModelHandler(
task="fill-mask"))
Args:
task (str or enum.Enum): task supported by HuggingFace Pipelines.
Accepts a string task or an enum.Enum from PipelineTask.
model (str): path to the pretrained *model-id* on Hugging Face Models Hub
to use custom model for the chosen task. If the `model` already defines
the task then no need to specify the `task` parameter.
Use the *model-id* string instead of an actual model here.
Model-specific kwargs for `from_pretrained(..., **model_kwargs)` can be
specified with `model_kwargs` using `load_pipeline_args`.
Example Usage::
model_handler = HuggingFacePipelineModelHandler(
task="text-generation", model="meta-llama/Llama-2-7b-hf",
load_pipeline_args={'model_kwargs':{'quantization_map':config}})
inference_fn: the inference function to use during RunInference.
Default is _default_pipeline_inference_fn.
load_pipeline_args (Dict[str, Any]): keyword arguments to provide load
options while loading pipelines from Hugging Face. Defaults to None.
inference_args (Dict[str, Any]): Non-batchable arguments
required as inputs to the model's inference function.
Defaults to None.
min_batch_size: the minimum batch size to use when batching inputs.
max_batch_size: the maximum batch size to use when batching inputs.
large_model: set to true if your model is large enough to run into
memory pressure if you load multiple copies. Given a model that
consumes N memory and a machine with W cores and M memory, you should
set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
**Supported Versions:** HuggingFacePipelineModelHandler supports
transformers>=4.18.0.
"""
self._task = task
self._model = model
self._inference_fn = inference_fn
self._load_pipeline_args = load_pipeline_args if load_pipeline_args else {}
self._inference_args = inference_args if inference_args else {}
self._batching_kwargs = {}
self._framework = "torch"
self._env_vars = kwargs.get('env_vars', {})
if min_batch_size is not None:
self._batching_kwargs['min_batch_size'] = min_batch_size
if max_batch_size is not None:
self._batching_kwargs['max_batch_size'] = max_batch_size
self._large_model = large_model
_validate_constructor_args_hf_pipeline(self._task, self._model)
[docs] def load_model(self):
"""Loads and initializes the pipeline for processing."""
return pipeline(
task=self._task, model=self._model, **self._load_pipeline_args)
[docs] def run_inference(
self,
batch: Sequence[str],
pipeline: Pipeline,
inference_args: Optional[Dict[str, Any]] = None
) -> Iterable[PredictionResult]:
"""
Runs inferences on a batch of examples passed as a string resource.
These can either be string sentences, or string path to images or
audio files.
Args:
batch: A sequence of strings resources.
pipeline: A Hugging Face Pipeline.
inference_args: Non-batchable arguments required as inputs to the model's
inference function.
Returns:
An Iterable of type PredictionResult.
"""
inference_args = {} if not inference_args else inference_args
predictions = self._inference_fn(batch, pipeline, inference_args)
return _convert_to_result(batch, predictions)
[docs] def update_model_path(self, model_path: Optional[str] = None):
"""
Updates the pretrained model used by the Hugging Face Pipeline task.
Make sure that the new model does the same task as initial model.
Args:
model_path (str): (Optional) Path to the new trained model
from Hugging Face. Defaults to None.
"""
self._model = model_path if model_path else self._model
[docs] def get_num_bytes(self, batch: Sequence[str]) -> int:
"""
Returns:
The number of bytes of input batch elements.
"""
return sum(sys.getsizeof(element) for element in batch)
[docs] def batch_elements_kwargs(self):
return self._batching_kwargs
[docs] def share_model_across_processes(self) -> bool:
return self._large_model
[docs] def get_metrics_namespace(self) -> str:
"""
Returns:
A namespace for metrics collected by the RunInference transform.
"""
return 'BeamML_HuggingFacePipelineModelHandler'