This notebook demonstrates how to run inference on your custom framework using the ModelHandler class.
Named-entity recognition (NER) is one of the most common tasks for natural language processing (NLP). NLP locates named entities in unstructured text and classifies the entities using pre-defined labels, such as person name, organization, date, and so on.
This example illustrates how to use the popular spaCy
package to load a machine learning (ML) model and perform inference in an Apache Beam pipeline using the RunInference PTransform
.
For more information about the RunInference API, see About Beam ML
in the Apache Beam documentation.
Install package dependencies
The RunInference library is available in Apache Beam versions 2.40 and later.
For this example, you need to install spaCy
and pandas
. A small NER model, en_core_web_sm
, is also installed, but you can use any valid spaCy
model.
# 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
Learn about spaCy
To learn more about spaCy
, create a spaCy
language object in memory using spaCy
's trained models.
You can install these models as Python packages.
For more information, see spaCy's Models and Languages
documentation.
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"
)
Create a model handler
This section demonstrates how to create your own ModelHandler
so that you can use spaCy
for inference.
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')])
Use KeyedModelHandler
to handle keyed data
If you have keyed data, use KeyedModelHandler
.
# 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
