Apache Beam RunInference with Hugging Face

Run in Google Colab View source on GitHub

This notebook shows how to use models from Hugging Face and Hugging Face pipeline in Apache Beam pipelines that uses the RunInference transform.

Apache Beam has built-in support for Hugging Face model handlers. Hugging Face has three model handlers:

For more information about using RunInference, see Get started with AI/ML pipelines in the Apache Beam documentation.

Install dependencies

Install both Apache Beam and the required dependencies for Hugging Face.

pip install torch --quiet
pip install tensorflow --quiet
pip install transformers==4.44.2 --quiet
pip install apache-beam[gcp]>=2.50 --quiet
 
from typing import Dict
from typing import Iterable
from typing import Tuple

import tensorflow as tf
import torch
from transformers import AutoTokenizer
from transformers import TFAutoModelForMaskedLM

import apache_beam as beam
from apache_beam.ml.inference.base import KeyedModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.inference.huggingface_inference import HuggingFacePipelineModelHandler
from apache_beam.ml.inference.huggingface_inference import HuggingFaceModelHandlerKeyedTensor
from apache_beam.ml.inference.huggingface_inference import HuggingFaceModelHandlerTensor
from apache_beam.ml.inference.huggingface_inference import PipelineTask

Use RunInference with Hugging Face pipelines

You can use Hugging Face pipelines with RunInference by using the HuggingFacePipelineModelHandler model handler. Similar to the Hugging Face pipelines, to instantiate the model handler, the model handler needs either the pipeline task or the model that defines the task. To pass any optional arguments to load the pipeline, use load_pipeline_args. To pass the optional arguments for inference, use inference_args.

You can define the pipeline task in one of the following two ways:

  • In the form of string, for example "translation". This option is similar to how the pipeline task is defined when using Hugging Face.
  • In the form of a PipelineTask enum object defined in Apache Beam, such as PipelineTask.Translation.

Create a model handler

This example demonstrates a task that translates text from English to Spanish.

model_handler = HuggingFacePipelineModelHandler(
    task=PipelineTask.Translation_XX_to_YY,
    model = "google/flan-t5-small",
    load_pipeline_args={'framework': 'pt'},
    inference_args={'max_length': 200}
)

Define the input examples

Use this code to define the input examples.

text = ["translate English to Spanish: How are you doing?",
        "translate English to Spanish: This is the Apache Beam project."]

Postprocess the results

The output from the RunInference transform is a PredictionResult object. Use that output to extract inferences, and then format and print the results.

class FormatOutput(beam.DoFn):
  """
  Extract the results from PredictionResult and print the results.
  """
  def process(self, element):
    example = element.example
    translated_text = element.inference[0]['translation_text']
    print(f'Example: {example}')
    print(f'Translated text: {translated_text}')
    print('-' * 80)

Run the pipeline

Use the following code to run the pipeline.

with beam.Pipeline() as beam_pipeline:
  examples = (
      beam_pipeline
      | "CreateExamples" >> beam.Create(text)
  )
  inferences = (
      examples
      | "RunInference" >> RunInference(model_handler)
      | "Print" >> beam.ParDo(FormatOutput())
  )

Example: translate English to Spanish: How are you doing?
Translated text: Cómo está acerca?
--------------------------------------------------------------------------------
Example: translate English to Spanish: This is the Apache Beam project.
Translated text: Esto es el proyecto Apache Beam.
--------------------------------------------------------------------------------

Try with a different model

One of the best parts of using RunInference is how easy it is to swap in different models. For example, if we wanted to use a larger model like DeepSeek-R1-Distill-Llama-8B outside of Colab (which has very tight memory constraints and limited GPU access), all we need to change is our ModelHandler:

model_handler = HuggingFacePipelineModelHandler(
    task=PipelineTask