How to Visualize Custom TFX Artifacts With InteractiveContext

August 14, 2022

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

Hejsan!

Resources

codesue