Support ModelCardGenerator in Vertex AI

tfx_addons/__init__.py

@@ -22,11 +22,17 @@
     "__version__",
 ] + list(_PKG_METADATA.keys())
 
+_ARTIFACTS = {
+    "ModelCard": "model_card_generator.artifact.ModelCard",
+}
+
 
 def __getattr__(name):  # pylint: disable=C0103
   # PEP-562: Lazy loaded attributes on python modules
   # NB(gcasassaez): We lazy load to avoid issues with dependencies not installed
   # for some subpackes
   if name in _ACTIVE_MODULES:
     return _importlib.import_module("." + name, __name__)
+  elif name in _ARTIFACTS:
+    return _importlib.import_module("." + _ARTIFACTS[name], __name__)
   raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

tfx_addons/model_card_generator/artifact.py

@@ -32,7 +32,7 @@ class ModelCard(Artifact):
     `<uri>/data/model_card.proto`.
     * the model card itself, located at the `<uri>/model_card/ directory`.
   """
-  TYPE_NAME = 'ModelCard'
+  TYPE_NAME = 'tfx_addons.ModelCard'
   TYPE_ANNOTATION = Metrics
 
 

tfx_addons/model_card_generator/artifact_test.py

@@ -39,8 +39,9 @@ def test_create_and_save_artifact(self):
                             self.store.get_artifacts_by_id([mc_artifact.id]))
     with self.subTest('properties'):
       with self.subTest('type_id'):
-        self.assertEqual(mc_artifact.type_id,
-                         self.store.get_artifact_type('ModelCard').id)
+        self.assertEqual(
+            mc_artifact.type_id,
+            self.store.get_artifact_type('tfx_addons.ModelCard').id)
       with self.subTest('uri'):
         self.assertEqual(mc_artifact.uri, '/path/to/model/card/assets')
       with self.subTest('name'):