fix: add param for multi-label per user's feedback (#887)

andrewferlitsch · web-flow · commit fda942ffbe00 · 2021-12-07T12:58:53.000-08:00
* fix: add param for multi-label per user's feedback

* fix: indentation

* test: update assert for new params

* lint: remove trailing whitespace
diff --git a/samples/model-builder/create_training_pipeline_image_classification_sample.py b/samples/model-builder/create_training_pipeline_image_classification_sample.py
@@ -24,6 +24,7 @@ def create_training_pipeline_image_classification_sample(
     display_name: str,
     dataset_id: int,
     model_display_name: Optional[str] = None,
+    multi_label: bool = False,
     training_fraction_split: float = 0.8,
     validation_fraction_split: float = 0.1,
     test_fraction_split: float = 0.1,
@@ -33,7 +34,11 @@ def create_training_pipeline_image_classification_sample(
 ):
     aiplatform.init(project=project, location=location)
 
-    job = aiplatform.AutoMLImageTrainingJob(display_name=display_name)
+    job = aiplatform.AutoMLImageTrainingJob(
+        display_name=display_name,
+        prediction_type='classification',
+        multi_label=multi_label
+    )
 
     my_image_ds = aiplatform.ImageDataset(dataset_id)
 
diff --git a/samples/model-builder/create_training_pipeline_image_classification_sample_test.py b/samples/model-builder/create_training_pipeline_image_classification_sample_test.py
@@ -44,7 +44,9 @@ def test_create_training_pipeline_image_classification_sample(
         project=constants.PROJECT, location=constants.LOCATION
     )
     mock_get_automl_image_training_job.assert_called_once_with(
-        display_name=constants.DISPLAY_NAME
+        display_name=constants.DISPLAY_NAME,
+        multi_label=False,
+        prediction_type='classification'
     )
     mock_run_automl_image_training_job.assert_called_once_with(
         dataset=mock_image_dataset,

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,9 @@ def test_create_training_pipeline_image_classification_sample(`
`44`	`44`	`project=constants.PROJECT, location=constants.LOCATION`
`45`	`45`	`)`
`46`	`46`	`mock_get_automl_image_training_job.assert_called_once_with(`
`47`		`- display_name=constants.DISPLAY_NAME`
	`47`	`+ display_name=constants.DISPLAY_NAME,`
	`48`	`+ multi_label=False,`
	`49`	`+ prediction_type='classification'`
`48`	`50`	`)`
`49`	`51`	`mock_run_automl_image_training_job.assert_called_once_with(`
`50`	`52`	`dataset=mock_image_dataset,`