diff --git a/tfx_addons/__init__.py b/tfx_addons/__init__.py index 90138fb0..9a8f6b65 100644 --- a/tfx_addons/__init__.py +++ b/tfx_addons/__init__.py @@ -22,6 +22,10 @@ "__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 @@ -29,4 +33,6 @@ def __getattr__(name): # pylint: disable=C0103 # 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}") diff --git a/tfx_addons/model_card_generator/artifact.py b/tfx_addons/model_card_generator/artifact.py index 0dbdccca..65f66bcd 100644 --- a/tfx_addons/model_card_generator/artifact.py +++ b/tfx_addons/model_card_generator/artifact.py @@ -32,7 +32,7 @@ class ModelCard(Artifact): `/data/model_card.proto`. * the model card itself, located at the `/model_card/ directory`. """ - TYPE_NAME = 'ModelCard' + TYPE_NAME = 'tfx_addons.ModelCard' TYPE_ANNOTATION = Metrics diff --git a/tfx_addons/model_card_generator/artifact_test.py b/tfx_addons/model_card_generator/artifact_test.py index 5ada6984..a53cc31d 100644 --- a/tfx_addons/model_card_generator/artifact_test.py +++ b/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'):