Bring your own ML model to Beam RunInference

# Uncomment the following lines to install the required packages.
# %pip install spacy pandas
# %pip install "apache-beam[gcp, dataframe, interactive]"
# !python -m spacy download en_core_web_sm

import spacy

nlp = spacy.load("en_core_web_sm")

# Add text strings.
text_strings = [
    "The New York Times is an American daily newspaper based in New York City with a worldwide readership.",
    "It was founded in 1851 by Henry Jarvis Raymond and George Jones, and was initially published by Raymond, Jones & Company."
]

# Check which entities spaCy can recognize.
doc = nlp(text_strings[0])

for ent in doc.ents:
    print(ent.text, ent.start_char, ent.end_char, ent.label_)

The New York Times 0 18 ORG
American 25 33 NORP
daily 34 39 DATE
New York City 59 72 GPE

# Visualize the results.
from spacy import displacy
displacy.render(doc, style="ent")

# Visualize another example.
displacy.render(nlp(text_strings[1]), style="ent")

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions

import warnings
warnings.filterwarnings("ignore")


pipeline = beam.Pipeline()

# Print the results for verification.
with pipeline as p:
    (p 
    | "CreateSentences" >> beam.Create(text_strings)
    | beam.Map(print)
    )

The New York Times is an American daily newspaper based in New York City with a worldwide readership.
It was founded in 1851 by Henry Jarvis Raymond and George Jones, and was initially published by Raymond, Jones & Company.

# Define `SpacyModelHandler` to load the model and perform the inference.

from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.inference.base import ModelHandler
from apache_beam.ml.inference.base import PredictionResult
from spacy import Language
from typing import Any
from typing import Dict
from typing import Iterable
from typing import Optional
from typing import Sequence

class SpacyModelHandler(ModelHandler[str,
                                     PredictionResult,
                                     Language]):
    def __init__(
        self,
        model_name: str = "en_core_web_sm",
    ):
        """ Implementation of the ModelHandler interface for spaCy using text as input.

        Example Usage::

          pcoll | RunInference(SpacyModelHandler())

        Args:
          model_name: The spaCy model name. Default is en_core_web_sm.
        """
        self._model_name = model_name
        self._env_vars = {}

    def load_model(self) -> Language:
        """Loads and initializes a model for processing."""
        return spacy.load(self._model_name)

    def run_inference(
        self,
        batch: Sequence[str],
        model: Language,
        inference_args: Optional[Dict[str, Any]] = None
    ) -> Iterable[PredictionResult]:
        """Runs inferences on a batch of text strings.

        Args:
          batch: A sequence of examples as text strings. 
          model: A spaCy language model
          inference_args: Any additional arguments for an inference.

        Returns:
          An Iterable of type PredictionResult.
        """
        # Loop each text string, and use a tuple to store the inference results.
        predictions = []
        for one_text in batch:
            doc = model(one_text)
            predictions.append(
                [(ent.text, ent.start_char, ent.end_char, ent.label_) for ent in doc.ents])
        return [PredictionResult(x, y) for x, y in zip(batch, predictions)]

# Verify that the inference results are correct.
with pipeline as p:
    (p 
    | "CreateSentences" >> beam.Create(text_strings)
    | "RunInferenceSpacy" >> RunInference(SpacyModelHandler("en_core_web_sm"))
    | beam.Map(print)
    )

The New York Times is an American daily newspaper based in New York City with a worldwide readership.
It was founded in 1851 by Henry Jarvis Raymond and George Jones, and was initially published by Raymond, Jones & Company.
PredictionResult(example='The New York Times is an American daily newspaper based in New York City with a worldwide readership.', inference=[('The New York Times', 0, 18, 'ORG'), ('American', 25, 33, 'NORP'), ('daily', 34, 39, 'DATE'), ('New York City', 59, 72, 'GPE')])
PredictionResult(example='It was founded in 1851 by Henry Jarvis Raymond and George Jones, and was initially published by Raymond, Jones & Company.', inference=[('1851', 18, 22, 'DATE'), ('Henry Jarvis', 26, 38, 'PERSON'), ('Raymond', 39, 46, 'PERSON'), ('George Jones', 51, 63, 'PERSON'), ('Raymond, Jones & Company', 96, 120, 'ORG')])

# You can use these text strings with keys to distinguish examples.
text_strings_with_keys = [
    ("example_0", "The New York Times is an American daily newspaper based in New York City with a worldwide readership."),
    ("example_1", "It was founded in 1851 by Henry Jarvis Raymond and George Jones, and was initially published by Raymond, Jones & Company.")
]

from apache_beam.runners.interactive.interactive_runner import InteractiveRunner
from apache_beam.ml.inference.base import KeyedModelHandler
from apache_beam.dataframe.convert import to_dataframe

pipeline = beam.Pipeline(InteractiveRunner())

keyed_spacy_model_handler = KeyedModelHandler(SpacyModelHandler("en_core_web_sm"))

# Verify that the inference results are correct.
with pipeline as p:
    results = (p 
    | "CreateSentences" >> beam.Create(text_strings_with_keys)
    | "RunInferenceSpacy" >> RunInference(keyed_spacy_model_handler)
    # Generate a schema suitable for conversion to a dataframe using Map to Row objects.
    | 'ToRows' >> beam.Map(lambda row: beam.Row(key=row[0], text=row[1][0], predictions=row[1][1]))
    )

# Convert the results to a pandas dataframe.
import apache_beam.runners.interactive.interactive_beam as ib

beam_df = to_dataframe(results)
df = ib.collect(beam_df)

df