The ability to visually inspect TFX artifacts with InteractiveContext is
invaluable for debugging TFX pipelines. However, until a couple months ago,
when I hacked on a library to display TFX artifacts with Streamlit,
I thought you could only use InteractiveContext to visualize TFMA and TFDV
artifacts. As it turns out, you can also use InteractiveContext to visualize
custom TFX artifacts. I’ll describe how to do so in this post.
Background
Before we get started, let’s walk through how InteractiveContext visualizes
standard artifacts.
When you call InteractiveContext.show() on an artifact, InteractiveContext checks
whether there’s a visualization for the artifact in its visualization registry.
Then, it calls the visualization’s display method if a visualization is available.
Here is the source code for InteractiveContext.show():
from tfx.orchestration.experimental.interactive import notebook_utils
from tfx.orchestration.experimental.interactive import visualizations
@notebook_utils.requires_ipython
def show(self, item: object) -> None:
"""Show the given object in an IPython notebook display."""
from IPython.core.display import display
from IPython.core.display import HTML
if isinstance(item, types.Channel):
channel = item
artifacts = channel.get()
for artifact in artifacts:
artifact_heading = 'Artifact at %s' % html.escape(artifact.uri)
display(HTML('<b>%s</b><br/><br/>' % artifact_heading))
visualization = visualizations.get_registry().get_visualization(
artifact.type_name)
if visualization:
visualization.display(artifact)
else:
display(item)
Default visualizations are defined in the standard_visualizations module, and they’re registered in the InteractiveContext’s constructor.
In this tutorial, we’ll use InteractiveContext to create and display a
greeting card.
Setup
Working in an interactive notebook, install TFX, verify the installation, then restart the notebook kernel.
! pip install tfx
import tfx
tfx.__version__
get_ipython().kernel.do_shutdown(restart=True)
Create a custom artifact, component, and visualizer
Here we create a custom Card artifact, a create_greeting_card component that
produces the artifact, and a CardVisualization that we can use to display the
artifact with InteractiveContext.
Due to serialization issues that arise when defining a Python function component in the same file as its pipeline, we write the component to a module.
The final result can be a lot to take in all at once, so let’s look at each step one at a time and then put everything together.
Import packages
Import the packages needed to create our component, artifact, and visualizer.
from pathlib import Path
from textwrap import dedent
from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import OutputArtifact, Parameter
from tfx.orchestration.experimental.interactive.visualizations import ArtifactVisualization
from tfx.types.artifact import Artifact
Create the Card artifact
class Card(Artifact):
TYPE_NAME = 'Card'
Create the create_greeting_card component
The create_greeting_card component takes a user-supplied greeting string, wraps
it in HTML, and saves the resulting HTML string to a file at the Card
artifact’s URI. (The style in the HTML string is adapted from a
CodePen by Álvaro.)
@component
def create_greeting_card(
greeting: Parameter[str],
card: OutputArtifact[Card],
):
card_template = dedent(
'''
<!DOCTYPE html>
<html>
<head>
<link rel="stylesheet" href="https://fonts.googleapis.com/css?family=Montserrat">
<style>
.container {{
display: grid;
place-items: center;
text-align: center;
}}
.greeting {{
--color-1: #DF8453;
--color-2: #3D8DAE;
--color-3: #E4A9A8;
animation: color-animation 4s linear infinite;
font-family: "Montserrat", sans-serif;
font-weight: 800;
font-size: 8.5vw;
text-transform: uppercase;
}}
@keyframes color-animation {{
0% {{color: var(--color-1)}}
32% {{color: var(--color-1)}}
33% {{color: var(--color-2)}}
65% {{color: var(--color-2)}}
66% {{color: var(—color-3)}}
99% {{color: var(—color-3)}}
100% {{color: var(—color-1)}}
}}
</style>
</head>
<body>
<div class="container">
<h2><span class="greeting">{greeting}</span></h2>
</div>
</body>
</html>
'''
)
p = Path(card.uri) / 'greeting.html'
p.write_text(card_template.format(greeting=greeting))
Create the CardVisualization class
To create the CardVisualization, we subclass the ArtifactVisualization
abstract base class. Then, we override the ARTIFACT_TYPE property with Card,
the type of artifact the visualization applies to. Next, we override the
display method to read an HTML file from the Card artifact’s URI and render
the resulting HTML string.
class CardVisualization(ArtifactVisualization):
ARTIFACT_TYPE = Card
def display(self, artifact: Artifact):
from IPython.display import display, HTML
files = Path(artifact.uri).glob('*.html')
card = next(files).read_text()
display(HTML(card))
Write the component to a module
Piece everything together and write the component to a module.
%%writefile custom_components.py
from pathlib import Path
from textwrap import dedent
from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import OutputArtifact, Parameter
from tfx.orchestration.experimental.interactive.visualizations import ArtifactVisualization
from tfx.types.artifact import Artifact
class Card(Artifact):
TYPE_NAME = 'Card'
@component
def create_greeting_card(
greeting: Parameter[str],
card: OutputArtifact[Card],
):
card_template = dedent(
'''
<!DOCTYPE html>
<html>
<head>
<link rel="stylesheet" href="https://fonts.googleapis.com/css?family=Montserrat">
<style>
.container {{
display: grid;
place-items: center;
text-align: center;
}}
.greeting {{
--color-1: #DF8453;
--color-2: #3D8DAE;
--color-3: #E4A9A8;
animation: color-animation 4s linear infinite;
font-family: "Montserrat", sans-serif;
font-weight: 800;
font-size: 8.5vw;
text-transform: uppercase;
}}
@keyframes color-animation {{
0% {{color: var(--color-1)}}
32% {{color: var(--color-1)}}
33% {{color: var(--color-2)}}
65% {{color: var(--color-2)}}
66% {{color: var(—color-3)}}
99% {{color: var(—color-3)}}
100% {{color: var(—color-1)}}
}}
</style>
</head>
<body>
<div class="container">
<h2><span class="greeting">{greeting}</span></h2>
</div>
</body>
</html>
'''
)
p = Path(card.uri) / 'greeting.html'
p.write_text(card_template.format(greeting=greeting))
class CardVisualization(ArtifactVisualization):
ARTIFACT_TYPE = Card
def display(self, artifact: Artifact):
from IPython.display import display, HTML
files = Path(artifact.uri).glob('*.html')
card = next(files).read_text()
display(HTML(card))
Visualize the custom artifact with InteractiveContext
Now, let’s visualize the custom artifact.
First, we import create_greeting_card, CardVisualization, and other necessary
packages.
from custom_components import create_greeting_card, CardVisualization
from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext
from tfx.orchestration.experimental.interactive import visualizations
Next, we create an InteractiveContext and add the CardVisualization to the
visualizations registry. This tells InteractiveContext how to visualize a Card
artifact when we call InteractiveContext.show() for a Card .
context = InteractiveContext()
visualizations.get_registry().register(CardVisualization)
Finally, we run the create_greeting_card component and display its Card.
greeting_card = create_greeting_card(greeting='Hejsan!')
context.run(greeting_card)
context.show(greeting_card.outputs['card'])
Artifact at /tmp/tfx-interactive-2022-08-14T03_10_45.839266-n50f0_ae/create_greeting_card/card/1
Hejsan!
This is a silly example. Would you really want to use a machine learning pipeline to create and display a greeting card? Probably not. But you can use similar logic to display model cards and data cards.
Summary
You now understand how InteractiveContext visualizes artifacts and how to create
custom artifact visualizations.
To create a custom visualization:
- Subclass the
ArtifactVisualizationabstract base class. - Override the
ARTIFACT_TYPEproperty with the type of artifact the visualization applies to. - Override the
displaymethod to read relevant content from the artifact’s URI and to render the content.
When you’re ready to display your artifact:
- Add the visualization to
InteractiveContext’s visualization registry. - Run
InteractiveContext.show()on the artifact.
Next Steps
As a next step, I challenge you to create an ExampleVisualization to display
decoded Example records. You may use the code from the
TFX Keras Component Tutorial
as a starting point.