def __init__(self,
model_export: channel.Channel,
model_blessing: channel.Channel,
push_destination: Optional[pusher_pb2.PushDestination] = None,
name: Text = None,
custom_config: Optional[Dict[Text, Any]] = None,
executor_class: Optional[Type[
base_executor.BaseExecutor]] = executor.Executor,
outputs: base_component.ComponentOutputs = None):
component_name = 'Pusher'
input_dict = {
'model_export': channel.as_channel(model_export),
'model_blessing': channel.as_channel(model_blessing),
}
exec_properties = {
'custom_config': custom_config,
}
if push_destination is None:
if executor_class == executor.Executor:
raise ValueError(
'push_destination is required unless custom '
'executor_class is supplied that does not require it.')
else:
exec_properties['push_destination'] = (
json_format.MessageToJson(push_destination))
super(Pusher, self).__init__(component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor_class,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
model_export,
model_blessing,
push_destination,
name = None,
custom_config = None,
executor_class = executor.Executor,
outputs = None):
component_name = 'Pusher'
input_dict = {
'model_export': channel.as_channel(model_export),
'model_blessing': channel.as_channel(model_blessing),
}
exec_properties = {
'push_destination': json_format.MessageToJson(push_destination),
'custom_config': custom_config,
}
super(Pusher, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor_class,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
examples: channel.Channel,
model: channel.Channel,
blessing: Optional[channel.Channel] = None,
name: Optional[Text] = None):
"""Construct a ModelValidator component.
Args:
examples: A Channel of 'ExamplesPath' type, usually produced by ExampleGen
component.
model: A Channel of 'ModelExportPath' type, usually produced by Trainer
component.
blessing: Optional output channel of 'ModelBlessingPath' for result of
blessing.
name: Optional unique name. Necessary if multiple ModelValidator
components are declared in the same pipeline.
"""
blessing = blessing or channel.Channel(
type_name='ModelBlessingPath',
artifacts=[types.Artifact('ModelBlessingPath')])
name = name or ''
spec = ModelValidatorSpec(
examples=channel.as_channel(examples),
model=channel.as_channel(model),
component_unique_name=name,
blessing=blessing)
super(ModelValidator, self).__init__(spec=spec, name=name)
def __init__(self,
transformed_examples: channel.Channel,
transform_output: channel.Channel,
schema: channel.Channel,
module_file: Text,
train_args: trainer_pb2.TrainArgs,
eval_args: trainer_pb2.EvalArgs,
custom_config: Optional[Dict[Text, Any]] = None,
name: Optional[Text] = None,
executor_class: Optional[Type[
base_executor.BaseExecutor]] = executor.Executor,
outputs: Optional[channel.Channel] = None):
component_name = 'Trainer'
input_dict = {
'transformed_examples': channel.as_channel(transformed_examples),
'transform_output': channel.as_channel(transform_output),
'schema': channel.as_channel(schema),
}
exec_properties = {
'train_args': json_format.MessageToJson(train_args),
'eval_args': json_format.MessageToJson(eval_args),
'module_file': module_file,
'custom_config': custom_config,
}
super(Trainer, self).__init__(component_name=component_name,
unique_name=name,
driver=driver.Driver,
executor=executor_class,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def test_construct(self):
examples = types.TfxType(type_name='ExamplesPath')
model_exports = types.TfxType(type_name='ModelExportPath')
evaluator = component.Evaluator(
examples=channel.as_channel([examples]),
model_exports=channel.as_channel([model_exports]))
self.assertEqual('ModelEvalPath', evaluator.outputs.output.type_name)
def __init__(self,
model_export: channel.Channel,
model_blessing: channel.Channel,
slack_token: Text,
channel_id: Text,
timeout_sec: int,
name: Optional[Text] = None,
outputs: Optional[base_component.ComponentOutputs] = None):
component_name = 'SlackComponent'
input_dict = {
'model_export': channel.as_channel(model_export),
'model_blessing': channel.as_channel(model_blessing),
}
exec_properties = {
'slack_token': slack_token,
'channel_id': channel_id,
'timeout_sec': timeout_sec,
}
super(SlackComponent, self).__init__(component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
model_export: channel.Channel,
model_blessing: channel.Channel,
push_destination: pusher_pb2.PushDestination,
name: Text = None,
custom_config: Optional[Dict[Text, Any]] = None,
executor_class: Optional[Type[
base_executor.BaseExecutor]] = executor.Executor,
outputs: base_component.ComponentOutputs = None):
component_name = 'Pusher'
input_dict = {
'model_export': channel.as_channel(model_export),
'model_blessing': channel.as_channel(model_blessing),
}
exec_properties = {
'push_destination': json_format.MessageToJson(push_destination),
'custom_config': custom_config,
}
super(Pusher, self).__init__(component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor_class,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
transformed_examples,
transform_output,
schema,
module_file,
train_args,
eval_args,
custom_config = None,
name = None,
outputs = None):
component_name = 'Trainer'
input_dict = {
'transformed_examples': channel.as_channel(transformed_examples),
'transform_output': channel.as_channel(transform_output),
'schema': channel.as_channel(schema),
}
exec_properties = {
'train_args': json_format.MessageToJson(train_args),
'eval_args': json_format.MessageToJson(eval_args),
'module_file': module_file,
'custom_config': custom_config,
}
super(Trainer, self).__init__(
component_name=component_name,
unique_name=name,
driver=driver.Driver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
transformed_examples,
transform_output,
schema,
module_file,
train_args,
eval_args,
custom_config = None,
name = None,
executor_class = executor.Executor,
outputs = None):
component_name = 'Trainer'
input_dict = {
'transformed_examples': channel.as_channel(transformed_examples),
'transform_output': channel.as_channel(transform_output),
'schema': channel.as_channel(schema),
}
exec_properties = {
'train_args': json_format.MessageToJson(train_args),
'eval_args': json_format.MessageToJson(eval_args),
'module_file': module_file,
'custom_config': custom_config,
}
super(Trainer, self).__init__(
component_name=component_name,
unique_name=name,
driver=driver.Driver,
executor=executor_class,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
examples: channel.Channel,
model_exports: channel.Channel,
feature_slicing_spec: Optional[
evaluator_pb2.FeatureSlicingSpec] = None,
output: Optional[channel.Channel] = None,
name: Optional[Text] = None):
"""Construct an Evaluator component.
Args:
examples: A Channel of 'ExamplesPath' type, usually produced by ExampleGen
component.
model_exports: A Channel of 'ModelExportPath' type, usually produced by
Trainer component.
feature_slicing_spec: Optional evaluator_pb2.FeatureSlicingSpec instance,
providing the way to slice the data.
output: Optional channel of 'ModelEvalPath' for result of evaluation.
name: Optional unique name. Necessary if multiple Evaluator components are
declared in the same pipeline.
"""
output = output or channel.Channel(
type_name='ModelEvalPath',
artifacts=[types.Artifact('ModelEvalPath')])
spec = EvaluatorSpec(
examples=channel.as_channel(examples),
model_exports=channel.as_channel(model_exports),
feature_slicing_spec=(feature_slicing_spec
or evaluator_pb2.FeatureSlicingSpec()),
output=output)
super(Evaluator, self).__init__(spec=spec, name=name)
def __init__(
self,
examples,
model_exports,
feature_slicing_spec = None,
name = None,
outputs = None):
component_name = 'Evaluator'
input_dict = {
'examples': channel.as_channel(examples),
'model_exports': channel.as_channel(model_exports),
}
exec_properties = {
'feature_slicing_spec':
json_format.MessageToJson(feature_slicing_spec or
evaluator_pb2.FeatureSlicingSpec()),
}
super(Evaluator, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
stats: channel.Channel,
schema: channel.Channel,
name: Text = None,
outputs: Dict[Text, channel.Channel] = None):
"""Construct an ExampleValidator component.
Args:
stats: A Channel of 'ExampleStatisticsPath' type. This should contain at
least 'eval' split. Other splits are ignored currently.
schema: A Channel of "SchemaPath' type.
name: Optional unique name. Necessary iff multiple ExampleValidator
components are declared in the same pipeline.
outputs: Optional dict from name to output channel.
"""
# TODO(zhitaoli): Move all constants to a constants.py.
component_name = 'ExampleValidator'
input_dict = {
'stats': channel.as_channel(stats),
'schema': channel.as_channel(schema)
}
exec_properties = {}
super(ExampleValidator, self).__init__(component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(
self,
examples: channel.Channel,
model_exports: channel.Channel,
feature_slicing_spec: Optional[evaluator_pb2.FeatureSlicingSpec] = None,
name: Optional[Text] = None,
outputs: Optional[base_component.ComponentOutputs] = None):
component_name = 'Evaluator'
input_dict = {
'examples': channel.as_channel(examples),
'model_exports': channel.as_channel(model_exports),
}
exec_properties = {
'feature_slicing_spec':
json_format.MessageToJson(feature_slicing_spec or
evaluator_pb2.FeatureSlicingSpec()),
}
super(Evaluator, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def test_construct(self):
examples = types.Artifact(type_name='ExamplesPath')
model_exports = types.Artifact(type_name='ModelExportPath')
evaluator = component.Evaluator(
examples=channel.as_channel([examples]),
model_exports=channel.as_channel([model_exports]))
self.assertEqual('ModelEvalPath', evaluator.outputs.output.type_name)
def test_construct(self):
examples = types.TfxArtifact(type_name='ExamplesPath')
model = types.TfxArtifact(type_name='ModelExportPath')
model_validator = component.ModelValidator(
examples=channel.as_channel([examples]),
model=channel.as_channel([model]))
self.assertEqual('ModelBlessingPath',
model_validator.outputs.blessing.type_name)
def test_construct(self):
example_vadalitor = component.ExampleValidator(
stats=channel.as_channel(
[types.TfxType(type_name='ExampleStatisticsPath', split='eval')]),
schema=channel.as_channel([types.TfxType(type_name='SchemaPath')]),
)
self.assertEqual('ExampleValidationPath',
example_vadalitor.outputs.output.type_name)
def testAsChannelDeprecated(self):
with mock.patch.object(logging, 'warning'):
warn_mock = mock.MagicMock()
logging.warning = warn_mock
channel.as_channel([standard_artifacts.Model()])
warn_mock.assert_called_once()
self.assertIn('tfx.utils.channel.as_channel has been renamed to',
warn_mock.call_args[0][5])
def test_construct(self):
example_vadalitor = component.ExampleValidator(
stats=channel.as_channel([
types.TfxType(type_name='ExampleStatisticsPath', split='eval')
]),
schema=channel.as_channel([types.TfxType(type_name='SchemaPath')]),
)
self.assertEqual('ExampleValidationPath',
example_vadalitor.outputs.output.type_name)
def test_construct_without_transform_output(self):
transformed_examples = types.TfxArtifact(type_name='ExamplesPath')
schema = types.TfxArtifact(type_name='SchemaPath')
trainer = component.Trainer(
module_file='/path/to/module/file',
examples=channel.as_channel([transformed_examples]),
schema=channel.as_channel([schema]),
train_args=trainer_pb2.TrainArgs(num_steps=100),
eval_args=trainer_pb2.EvalArgs(num_steps=50))
self.assertEqual('ModelExportPath', trainer.outputs.output.type_name)
def test_construct(self):
model_export = types.TfxArtifact(type_name='ModelExportPath')
model_blessing = types.TfxArtifact(type_name='ModelBlessingPath')
pusher = component.Pusher(
model_export=channel.as_channel([model_export]),
model_blessing=channel.as_channel([model_blessing]),
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='push_destination')))
self.assertEqual('ModelPushPath', pusher.outputs.model_push.type_name)
def test_construct_with_slice_spec(self):
examples = types.Artifact(type_name='ExamplesPath')
model_exports = types.Artifact(type_name='ModelExportPath')
evaluator = component.Evaluator(
examples=channel.as_channel([examples]),
model_exports=channel.as_channel([model_exports]),
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
self.assertEqual('ModelEvalPath', evaluator.outputs.output.type_name)
def test_construct_with_slice_spec(self):
examples = types.TfxType(type_name='ExamplesPath')
model_exports = types.TfxType(type_name='ModelExportPath')
evaluator = component.Evaluator(
examples=channel.as_channel([examples]),
model_exports=channel.as_channel([model_exports]),
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
self.assertEqual('ModelEvalPath', evaluator.outputs.output.type_name)
def test_construct(self):
transformed_examples = types.TfxType(type_name='ExamplesPath')
transform_output = types.TfxType(type_name='TransformPath')
schema = types.TfxType(type_name='SchemaPath')
trainer = component.Trainer(
module_file='/path/to/module/file',
transformed_examples=channel.as_channel([transformed_examples]),
transform_output=channel.as_channel([transform_output]),
schema=channel.as_channel([schema]),
train_args=trainer_pb2.TrainArgs(num_steps=100),
eval_args=trainer_pb2.EvalArgs(num_steps=50))
self.assertEqual('ModelExportPath', trainer.outputs.output.type_name)
def __init__(
self,
transformed_examples: channel.Channel = None,
transform_output: channel.Channel = None,
schema: channel.Channel = None,
module_file: Text = None,
train_args: trainer_pb2.TrainArgs = None,
eval_args: trainer_pb2.EvalArgs = None,
custom_config: Optional[Dict[Text, Any]] = None,
executor_class: Optional[Type[base_executor.BaseExecutor]] = None,
output: Optional[channel.Channel] = None,
name: Optional[Text] = None):
"""Construct a Trainer component.
Args:
transformed_examples: A Channel of 'ExamplesPath' type, serving as the
source of transformed examples. This is usually an output of Transform.
transform_output: A Channel of 'TransformPath' type, serving as the input
transform graph.
schema: A Channel of 'SchemaPath' type, serving as the schema of training
and eval data.
module_file: A python module file containing UDF model definition.
train_args: A trainer_pb2.TrainArgs instance, containing args used for
training. Current only num_steps is available.
eval_args: A trainer_pb2.EvalArgs instance, containing args used for
eval. Current only num_steps is available.
custom_config: A dict which contains the training job parameters to be
passed to Google Cloud ML Engine. For the full set of parameters
supported by Google Cloud ML Engine, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#Job
executor_class: Optional custom executor class.
output: Optional 'ModelExportPath' channel for result of exported models.
name: Optional unique name. Necessary iff multiple Trainer components are
declared in the same pipeline.
"""
output = output or channel.Channel(
type_name='ModelExportPath',
artifacts=[types.TfxArtifact('ModelExportPath')])
spec = TrainerSpec(
transformed_examples=channel.as_channel(transformed_examples),
transform_output=channel.as_channel(transform_output),
schema=channel.as_channel(schema),
train_args=train_args,
eval_args=eval_args,
module_file=module_file,
custom_config=custom_config,
output=output)
super(Trainer, self).__init__(spec=spec,
custom_executor_class=executor_class,
name=name)
def test_channel_as_channel_success(self):
instance_a = types.TfxType('MyTypeName')
instance_b = types.TfxType('MyTypeName')
chnl_original = channel.Channel(
'MyTypeName', static_artifact_collection=[instance_a, instance_b])
chnl_result = channel.as_channel(chnl_original)
self.assertEqual(chnl_original, chnl_result)
def test_construct(self):
train_examples = types.TfxType(type_name='ExamplesPath', split='train')
eval_examples = types.TfxType(type_name='ExamplesPath', split='eval')
statistics_gen = component.StatisticsGen(
input_data=channel.as_channel([train_examples, eval_examples]))
self.assertEqual('ExampleStatisticsPath',
statistics_gen.outputs.output.type_name)
def __init__(self,
input_config: example_gen_pb2.Input,
output_config: Optional[example_gen_pb2.Output] = None,
component_name: Optional[Text] = 'ExampleGen',
example_artifacts: Optional[channel.Channel] = None,
name: Optional[Text] = None):
"""Construct an QueryBasedExampleGen component.
Args:
input_config: An example_gen_pb2.Input instance, providing input
configuration.
output_config: An example_gen_pb2.Output instance, providing output
configuration. If unset, default splits will be 'train' and 'eval' with
size 2:1.
component_name: Name of the component, should be unique per component
class. Default to 'ExampleGen', can be overwritten by sub-classes.
example_artifacts: Optional channel of 'ExamplesPath' for output train and
eval examples.
name: Unique name for every component class instance.
"""
# Configure outputs.
output_config = output_config or utils.make_default_output_config(
input_config)
example_artifacts = example_artifacts or channel.as_channel([
types.TfxArtifact('ExamplesPath', split=split_name)
for split_name in utils.generate_output_split_names(
input_config, output_config)
])
spec = QueryBasedExampleGenSpec(component_name=component_name,
input_config=input_config,
output_config=output_config,
examples=example_artifacts)
super(_QueryBasedExampleGen, self).__init__(spec=spec, name=name)
def test_construct(self):
train_examples = types.TfxType(type_name='ExamplesPath', split='train')
eval_examples = types.TfxType(type_name='ExamplesPath', split='eval')
statistics_gen = component.StatisticsGen(
input_data=channel.as_channel([train_examples, eval_examples]))
self.assertEqual('ExampleStatisticsPath',
statistics_gen.outputs.output.type_name)
def test_channel_as_channel_success(self):
instance_a = types.Artifact('MyTypeName')
instance_b = types.Artifact('MyTypeName')
chnl_original = channel.Channel('MyTypeName',
artifacts=[instance_a, instance_b])
chnl_result = channel.as_channel(chnl_original)
self.assertEqual(chnl_original, chnl_result)
def __init__(self,
executor: Any,
input_base: Optional[channel.Channel] = None,
input_config: Optional[example_gen_pb2.Input] = None,
output_config: Optional[example_gen_pb2.Output] = None,
component_name: Optional[Text] = 'ExampleGen',
unique_name: Optional[Text] = None,
outputs: Optional[base_component.ComponentOutputs] = None):
if input_base is None and input_config is None:
raise RuntimeError(
'One of input_base and input_config must be set.')
input_dict = {
'input-base': channel.as_channel(input_base)
} if input_base else {}
# Default value need to be set in component instead of executor as output
# artifacts depend on it.
self._input_config = input_config or utils.make_default_input_config()
self._output_config = output_config or utils.make_default_output_config(
self._input_config)
exec_properties = {
'input': json_format.MessageToJson(self._input_config),
'output': json_format.MessageToJson(self._output_config)
}
super(ExampleGen, self).__init__(
component_name=component_name,
unique_name=unique_name,
driver=driver.Driver if input_base else base_driver.BaseDriver,
executor=executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def test_construct(self):
source_data_dir = os.path.join(
os.path.dirname(__file__), 'testdata', 'taxi')
preprocessing_fn_file = os.path.join(source_data_dir, 'module',
'preprocess.py')
transform = component.Transform(
input_data=channel.as_channel([
types.TfxType(type_name='ExamplesPath', split='train'),
types.TfxType(type_name='ExamplesPath', split='eval'),
]),
schema=channel.as_channel([types.TfxType(type_name='SchemaPath')]),
module_file=preprocessing_fn_file,
)
self.assertEqual('TransformPath',
transform.outputs.transform_output.type_name)
self.assertEqual('ExamplesPath',
transform.outputs.transformed_examples.type_name)
def test_construct(self):
source_data_dir = os.path.join(os.path.dirname(__file__), 'testdata',
'taxi')
preprocessing_fn_file = os.path.join(source_data_dir, 'module',
'preprocess.py')
transform = component.Transform(
input_data=channel.as_channel([
types.TfxType(type_name='ExamplesPath', split='train'),
types.TfxType(type_name='ExamplesPath', split='eval'),
]),
schema=channel.as_channel([types.TfxType(type_name='SchemaPath')]),
module_file=preprocessing_fn_file,
)
self.assertEqual('TransformPath',
transform.outputs.transform_output.type_name)
self.assertEqual('ExamplesPath',
transform.outputs.transformed_examples.type_name)
def __init__(self, examples, model, name=None, outputs=None):
component_name = 'ModelValidator'
input_dict = {
'examples': channel.as_channel(examples),
'model': channel.as_channel(model),
}
exec_properties = {
'blessed_model': None, # This will be set in driver.
'blessed_model_id': None, # This will be set in driver.
}
super(ModelValidator, self).__init__(component_name=component_name,
unique_name=name,
driver=driver.Driver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def test_construct_subclass(self):
input_base = types.TfxArtifact(type_name='ExternalPath')
example_gen = TestExampleGenComponent(
input_base=channel.as_channel([input_base]))
self.assertEqual('ExamplesPath', example_gen.outputs.examples.type_name)
artifact_collection = example_gen.outputs.examples.get()
self.assertEqual('train', artifact_collection[0].split)
self.assertEqual('eval', artifact_collection[1].split)
def test_construct(self):
input_base = types.Artifact(type_name='ExternalPath')
import_example_gen = component.ImportExampleGen(
input_base=channel.as_channel([input_base]))
self.assertEqual('ExamplesPath',
import_example_gen.outputs.examples.type_name)
artifact_collection = import_example_gen.outputs.examples.get()
self.assertEqual('train', artifact_collection[0].split)
self.assertEqual('eval', artifact_collection[1].split)
def __init__(
self,
model_export: channel.Channel,
model_blessing: channel.Channel,
push_destination: Optional[pusher_pb2.PushDestination] = None,
custom_config: Optional[Dict[Text, Any]] = None,
executor_class: Optional[Type[base_executor.BaseExecutor]] = None,
model_push: Optional[channel.Channel] = None,
name: Optional[Text] = None):
"""Construct a Pusher component.
Args:
model_export: A Channel of 'ModelExportPath' type, usually produced by
Trainer component.
model_blessing: A Channel of 'ModelBlessingPath' type, usually produced by
ModelValidator component.
push_destination: A pusher_pb2.PushDestination instance, providing
info for tensorflow serving to load models. Optional if executor_class
doesn't require push_destination.
custom_config: A dict which contains the deployment job parameters to be
passed to Google Cloud ML Engine. For the full set of parameters
supported by Google Cloud ML Engine, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models
executor_class: Optional custom python executor class.
model_push: Optional output 'ModelPushPath' channel with result of push.
name: Optional unique name. Necessary if multiple Pusher components are
declared in the same pipeline.
"""
model_push = model_push or channel.Channel(
type_name='ModelPushPath',
artifacts=[types.Artifact('ModelPushPath')])
if push_destination is None and not executor_class:
raise ValueError(
'push_destination is required unless a custom '
'executor_class is supplied that does not require '
'it.')
spec = PusherSpec(model_export=channel.as_channel(model_export),
model_blessing=channel.as_channel(model_blessing),
push_destination=push_destination,
custom_config=custom_config,
model_push=model_push)
super(Pusher, self).__init__(spec=spec,
custom_executor_class=executor_class,
name=name)
def test_construct_custom_executor(self):
input_base = types.TfxArtifact(type_name='ExternalPath')
example_gen = component.FileBasedExampleGen(
input_base=channel.as_channel([input_base]),
executor_class=TestExampleGenExecutor)
self.assertEqual(driver.Driver, example_gen.driver_class)
self.assertEqual('ExamplesPath', example_gen.outputs.examples.type_name)
artifact_collection = example_gen.outputs.examples.get()
self.assertEqual('train', artifact_collection[0].split)
self.assertEqual('eval', artifact_collection[1].split)
def test_construct_subclass_file_based(self):
input_base = types.TfxArtifact(type_name='ExternalPath')
example_gen = TestFileBasedExampleGenComponent(
input_base=channel.as_channel([input_base]))
self.assertIn('input_base', example_gen.inputs.get_all())
self.assertEqual(driver.Driver, example_gen.driver_class)
self.assertEqual('ExamplesPath', example_gen.outputs.examples.type_name)
artifact_collection = example_gen.outputs.examples.get()
self.assertEqual('train', artifact_collection[0].split)
self.assertEqual('eval', artifact_collection[1].split)
def __init__(self,
input_data,
schema,
module_file,
name = None,
outputs = None):
component_name = 'Transform'
input_dict = {
'input_data': channel.as_channel(input_data),
'schema': channel.as_channel(schema)
}
exec_properties = {
'module_file': module_file,
}
super(Transform, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
examples,
model,
name = None,
outputs = None):
component_name = 'ModelValidator'
input_dict = {
'examples': channel.as_channel(examples),
'model': channel.as_channel(model),
}
exec_properties = {
'blessed_model': None, # This will be set in driver.
'blessed_model_id': None, # This will be set in driver.
}
super(ModelValidator, self).__init__(
component_name=component_name,
unique_name=name,
driver=driver.Driver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
model_export,
model_blessing,
push_destination,
name = None,
custom_config = None,
outputs = None):
component_name = 'Pusher'
input_dict = {
'model_export': channel.as_channel(model_export),
'model_blessing': channel.as_channel(model_blessing),
}
exec_properties = {
'push_destination': json_format.MessageToJson(push_destination),
'custom_config': custom_config,
}
super(Pusher, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
input_base,
name = None,
outputs = None):
component_name = 'CsvExampleGen'
input_dict = {'input-base': channel.as_channel(input_base)}
exec_properties = {}
super(CsvExampleGen, self).__init__(
component_name=component_name,
unique_name=name,
driver=driver.Driver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
input_data,
name = None,
outputs = None):
"""Constructs a StatisticsGen component.
Args:
input_data: A Channel of 'ExamplesPath' type. This should contain two
splits 'train' and 'eval'.
name: Optional unique name. Necessary iff multiple StatisticsGen
components are declared in the same pipeline.
outputs: Optional dict from name to output channel.
"""
component_name = 'StatisticsGen'
input_dict = {'input_data': channel.as_channel(input_data)}
exec_properties = {}
super(StatisticsGen, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def __init__(self,
stats,
name = None,
outputs = None):
"""Constructs a SchemaGen component.
Args:
stats: A Channel of 'ExampleStatisticsPath' type. This should contain at
least 'train' split. Other splits are ignored currently.
name: Optional unique name. Necessary iff multiple SchemaGen components
are declared in the same pipeline.
outputs: Optional dict from name to output channel.
"""
component_name = 'SchemaGen'
input_dict = {'stats': channel.as_channel(stats)}
exec_properties = {}
super(SchemaGen, self).__init__(
component_name=component_name,
unique_name=name,
driver=base_driver.BaseDriver,
executor=executor.Executor,
input_dict=input_dict,
outputs=outputs,
exec_properties=exec_properties)
def test_construct(self):
input_base = types.TfxType(type_name='ExternalPath')
csv_example_gen = component.CsvExampleGen(
input_base=channel.as_channel([input_base]))
self.assertEqual('ExamplesPath', csv_example_gen.outputs.examples.type_name)
def test_construct(self):
schema_gen = component.SchemaGen(
stats=channel.as_channel(
[types.TfxType(type_name='ExampleStatisticsPath', split='train')]))
self.assertEqual('SchemaPath', schema_gen.outputs.output.type_name)