# Path: model_card_toolkit/tfx/artifact.py
# class ModelCard(Artifact):
#   TYPE_NAME = 'ModelCard'
#   TYPE_ANNOTATION = Metrics
# def create_and_save_artifact(
#     artifact_name: str, artifact_uri: str,
#     store: mlmd.MetadataStore) -> metadata_store_pb2.Artifact:
# 
# Path: model_card_toolkit/tfx/executor.py
# _DEFAULT_MODEL_CARD_FILE_NAME = 'model_card.html'
# class Executor(BaseExecutor):
#   def _tfma_source(
#       self, input_dict: Dict[str,
#                              List[types.Artifact]]) -> Optional[src.TfmaSource]:
#   def _tfdv_source(
#       self, input_dict: Dict[str,
#                              List[types.Artifact]]) -> Optional[src.TfdvSource]:
#   def _model_source(
#       self,
#       input_dict: Dict[str, List[types.Artifact]]) -> Optional[src.ModelSource]:
#   def Do(self, input_dict: Dict[str, List[types.Artifact]],
#          output_dict: Dict[str, List[types.Artifact]],
#          exec_properties: Dict[str, Any]) -> None:


from typing import Any, List, Tuple, Optional
from model_card_toolkit.tfx import artifact
from model_card_toolkit.tfx import executor
from tfx import types
from tfx.dsl.components.base import executor_spec
from tfx.dsl.components.base.base_component import BaseComponent
from tfx.types import component_spec
from tfx.types import standard_artifacts
from tfx.types import standard_component_specs

MODEL_CARD_KEY = 'model_card'


class ModelCardGeneratorSpec(component_spec.ComponentSpec):
  """Component spec for the ModelCardGenerator."""
  PARAMETERS = {
      'json':
          component_spec.ExecutionParameter(type=str, optional=True),
      # template_io's type is List[Tuple[str, str]],
      # but we need List[Any] to pass ExecutionParameter.type_check().
      # See below link for details.
      # https://github.com/tensorflow/tfx/blob/4ff5e97b09540ff8a858076a163ecdf209716324/tfx/types/component_spec.py#L308
      'template_io':
          component_spec.ExecutionParameter(
              type=List[Any], optional=True)
  }
  INPUTS = {
      standard_component_specs.STATISTICS_KEY:
          component_spec.ChannelParameter(
              type=standard_artifacts.ExampleStatistics, optional=True),
      standard_component_specs.EVALUATION_KEY:
          component_spec.ChannelParameter(
              type=standard_artifacts.ModelEvaluation, optional=True),
      standard_component_specs.PUSHED_MODEL_KEY:
          component_spec.ChannelParameter(
              type=standard_artifacts.PushedModel, optional=True),
  }
  OUTPUTS = {
      MODEL_CARD_KEY:

# Path: model_card_toolkit/tfx/artifact.py
# class ModelCard(Artifact):
#   TYPE_NAME = 'ModelCard'
#   TYPE_ANNOTATION = Metrics
# def create_and_save_artifact(
#     artifact_name: str, artifact_uri: str,
#     store: mlmd.MetadataStore) -> metadata_store_pb2.Artifact:
# 
# Path: model_card_toolkit/tfx/executor.py
# _DEFAULT_MODEL_CARD_FILE_NAME = 'model_card.html'
# class Executor(BaseExecutor):
#   def _tfma_source(
#       self, input_dict: Dict[str,
#                              List[types.Artifact]]) -> Optional[src.TfmaSource]:
#   def _tfdv_source(
#       self, input_dict: Dict[str,
#                              List[types.Artifact]]) -> Optional[src.TfdvSource]:
#   def _model_source(
#       self,
#       input_dict: Dict[str, List[types.Artifact]]) -> Optional[src.ModelSource]:
#   def Do(self, input_dict: Dict[str, List[types.Artifact]],
#          output_dict: Dict[str, List[types.Artifact]],
#          exec_properties: Dict[str, Any]) -> None:


# Path: model_card_toolkit/tfx/component.py
from typing import Any, List, Tuple, Optional
from model_card_toolkit.tfx import artifact
from model_card_toolkit.tfx import executor
from tfx import types
from tfx.dsl.components.base import executor_spec
from tfx.dsl.components.base.base_component import BaseComponent
from tfx.types import component_spec
from tfx.types import standard_artifacts
from tfx.types import standard_component_specs


MODEL_CARD_KEY = 'model_card'


class ModelCardGeneratorSpec(component_spec.ComponentSpec):
  """Component spec for the ModelCardGenerator."""
  PARAMETERS = {
      'json':
          component_spec.ExecutionParameter(type=str, optional=True),
      # template_io's type is List[Tuple[str, str]],
      # but we need List[Any] to pass ExecutionParameter.type_check().
      # See below link for details.
      # https://github.com/tensorflow/tfx/blob/4ff5e97b09540ff8a858076a163ecdf209716324/tfx/types/component_spec.py#L308
      'template_io':
          component_spec.ExecutionParameter(
              type=List[Any], optional=True)
  }
  INPUTS = {
      standard_component_specs.STATISTICS_KEY:
          component_spec.ChannelParameter(
              type=standard_artifacts.ExampleStatistics, optional=True),
      standard_component_specs.EVALUATION_KEY:
          component_spec.ChannelParameter(
              type=standard_artifacts.ModelEvaluation, optional=True),
      standard_component_specs.PUSHED_MODEL_KEY:
          component_spec.ChannelParameter(
              type=standard_artifacts.PushedModel, optional=True),
  }
  OUTPUTS = {
      MODEL_CARD_KEY:

          component_spec.ChannelParameter(type=artifact.ModelCard),