class SimpleComponentSpec(base_component.ComponentSpec):
PARAMETERS = {
'x': ExecutionParameter(type=int),
'y': ExecutionParameter(type=int, optional=True),
}
INPUTS = {'z': ChannelParameter(type_name='Z')}
OUTPUTS = {}
class QueryBasedExampleGenSpec(base_component.ComponentSpec):
"""Query-based ExampleGen component spec."""
PARAMETERS = {
'input_config': ExecutionParameter(type=example_gen_pb2.Input),
'output_config': ExecutionParameter(type=example_gen_pb2.Output),
}
INPUTS = {}
OUTPUTS = {
'examples': ChannelParameter(type_name='ExamplesPath'),
}
class _BasicComponentSpec(base_component.ComponentSpec):
PARAMETERS = {
'folds': ExecutionParameter(type=int),
'proto': ExecutionParameter(type=example_gen_pb2.Input, optional=True),
}
INPUTS = {
'input': ChannelParameter(type_name='InputType'),
}
OUTPUTS = {
'output': ChannelParameter(type_name='OutputType'),
}
class FileBasedExampleGenSpec(base_component.ComponentSpec):
"""File-based ExampleGen component spec."""
COMPONENT_NAME = 'ExampleGen'
PARAMETERS = {
'input_config': ExecutionParameter(type=example_gen_pb2.Input),
'output_config': ExecutionParameter(type=example_gen_pb2.Output),
}
INPUTS = {
'input_base': ChannelParameter(type_name='ExternalPath'),
}
OUTPUTS = {
'examples': ChannelParameter(type_name='ExamplesPath'),
}
class SlackComponentSpec(base_component.ComponentSpec):
"""ComponentSpec for Custom TFX Slack Component."""
PARAMETERS = {
'slack_token': ExecutionParameter(type=Text),
'channel_id': ExecutionParameter(type=Text),
'timeout_sec': ExecutionParameter(type=int),
}
INPUTS = {
'model_export': ChannelParameter(type_name='ModelExportPath'),
'model_blessing': ChannelParameter(type_name='ModelBlessingPath'),
}
OUTPUTS = {
'slack_blessing': ChannelParameter(type_name='ModelBlessingPath'),
}
class PusherSpec(base_component.ComponentSpec):
"""Pusher component spec."""
PARAMETERS = {
'push_destination':
ExecutionParameter(type=pusher_pb2.PushDestination, optional=True),
'custom_config':
ExecutionParameter(type=Dict[Text, Any], optional=True),
}
INPUTS = {
'model_export': ChannelParameter(type_name='ModelExportPath'),
'model_blessing': ChannelParameter(type_name='ModelBlessingPath'),
}
OUTPUTS = {
'model_push': ChannelParameter(type_name='ModelPushPath'),
}
class TrainerSpec(base_component.ComponentSpec):
"""Trainer component spec."""
COMPONENT_NAME = 'Trainer'
PARAMETERS = {
'train_args': ExecutionParameter(type=trainer_pb2.TrainArgs),
'eval_args': ExecutionParameter(type=trainer_pb2.EvalArgs),
'module_file': ExecutionParameter(type=(str, Text)),
'custom_config': ExecutionParameter(type=Dict[Text, Any],
optional=True),
}
INPUTS = {
'transformed_examples': ChannelParameter(type_name='ExamplesPath'),
'transform_output': ChannelParameter(type_name='TransformPath'),
'schema': ChannelParameter(type_name='SchemaPath'),
}
OUTPUTS = {'output': ChannelParameter(type_name='ModelExportPath')}
class TransformSpec(base_component.ComponentSpec):
"""Transform component spec."""
PARAMETERS = {
'module_file': ExecutionParameter(type=(str, Text), optional=True),
'preprocessing_fn': ExecutionParameter(type=(str, Text),
optional=True),
}
INPUTS = {
'input_data': ChannelParameter(type=standard_artifacts.Examples),
'schema': ChannelParameter(type=standard_artifacts.Schema),
}
OUTPUTS = {
'transform_output':
ChannelParameter(type=standard_artifacts.TransformResult),
'transformed_examples':
ChannelParameter(type=standard_artifacts.Examples),
}
class CustomHeadComponentSpec(base_component.ComponentSpec):
"""Example for a ComponentSpec of a Custom TFX Head Component."""
COMPONENT_NAME = 'Custom_Head_Component'
PARAMETERS = {
'string_execution_parameter': ExecutionParameter(type=Text),
'integer_execution_parameter': ExecutionParameter(type=int),
# don't change these two:
'input_config': ExecutionParameter(type=example_gen_pb2.Input),
'output_config': ExecutionParameter(type=example_gen_pb2.Output)
}
INPUTS = {
'input_example': ChannelParameter(type_name='RandomTypeNameForInput')
}
OUTPUTS = {
'output_example': ChannelParameter(type_name='RandomTypeNameForOutput')
}
class TrainerSpec(base_component.ComponentSpec):
"""Trainer component spec."""
PARAMETERS = {
'train_args': ExecutionParameter(type=trainer_pb2.TrainArgs),
'eval_args': ExecutionParameter(type=trainer_pb2.EvalArgs),
'module_file': ExecutionParameter(type=(str, Text), optional=True),
'trainer_fn': ExecutionParameter(type=(str, Text), optional=True),
'custom_config': ExecutionParameter(type=Dict[Text, Any],
optional=True),
}
INPUTS = {
'examples':
ChannelParameter(type=standard_artifacts.Examples),
'transform_output':
ChannelParameter(type=standard_artifacts.TransformResult,
optional=True),
'schema':
ChannelParameter(type=standard_artifacts.Schema),
}
OUTPUTS = {'output': ChannelParameter(type=standard_artifacts.Model)}
class ModelValidatorSpec(base_component.ComponentSpec):
"""ModelValidator component spec."""
PARAMETERS = {
'component_unique_name': ExecutionParameter(type=str),
}
INPUTS = {
'examples': ChannelParameter(type_name='ExamplesPath'),
'model': ChannelParameter(type_name='ModelExportPath'),
}
OUTPUTS = {
'blessing': ChannelParameter(type_name='ModelBlessingPath'),
}
class TransformSpec(base_component.ComponentSpec):
"""Transform component spec."""
PARAMETERS = {
'module_file': ExecutionParameter(type=(str, Text)),
}
INPUTS = {
'input_data': ChannelParameter(type_name='ExamplesPath'),
'schema': ChannelParameter(type_name='SchemaPath'),
}
OUTPUTS = {
'transform_output': ChannelParameter(type_name='TransformPath'),
'transformed_examples': ChannelParameter(type_name='ExamplesPath'),
}
class EvaluatorSpec(base_component.ComponentSpec):
"""Evaluator component spec."""
PARAMETERS = {
'feature_slicing_spec': ExecutionParameter(
type=evaluator_pb2.FeatureSlicingSpec),
}
INPUTS = {
'examples': ChannelParameter(type=standard_artifacts.Examples),
'model_exports': ChannelParameter(type=standard_artifacts.Model),
}
OUTPUTS = {
'output': ChannelParameter(type=standard_artifacts.ModelEvalResult),
}
class EvaluatorSpec(base_component.ComponentSpec):
"""Evaluator component spec."""
PARAMETERS = {
'feature_slicing_spec':
ExecutionParameter(type=evaluator_pb2.FeatureSlicingSpec),
}
INPUTS = {
'examples': ChannelParameter(type_name='ExamplesPath'),
'model_exports': ChannelParameter(type_name='ModelExportPath'),
}
OUTPUTS = {
'output': ChannelParameter(type_name='ModelEvalPath'),
}
class WrongTypeComponentSpecC(base_component.ComponentSpec):
COMPONENT_NAME = 'WrongTypeComponentC'
PARAMETERS = {}
INPUTS = {'x': ExecutionParameter(type=int)}
OUTPUTS = {}
class DuplicatePropertyComponentSpec(base_component.ComponentSpec):
PARAMETERS = {'x': ExecutionParameter(type=int)}
INPUTS = {'x': ChannelParameter(type_name='X')}
OUTPUTS = {}
class WrongTypeComponentSpecD(base_component.ComponentSpec):
PARAMETERS = {}
INPUTS = {'x': ExecutionParameter(type=int)}
OUTPUTS = {}