codesue

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.

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:

When you're ready to display your 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.

Resources