How to Visualize Custom TFX Artifacts With InteractiveContext
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.
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
.
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.)
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.
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.
%%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>
body {{
border: 1px solid #e6e9ef;
border-radius: 0.75rem;
}}
.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
and CardVisualization
from the module we wrote them to and create an InteractiveContext
. Next, we add CardVisualization
to the visualizations registry. This tells InteractiveContext
how to visualize a Card
artifact when we call InteractiveContext.show()
for a Card
. Finally, we run the create_greeting_card
component and display its Card
.
from custom_components import create_greeting_card, CardVisualization
from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext
from tfx.orchestration.experimental.interactive import visualizations
context = InteractiveContext()
visualizations.get_registry().register(CardVisualization)
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