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

* fix: add param for multi-label per user's feedback * fix: indentation * test: update assert for new params * lint: remove trailing whitespace
googleapis · Dec 7, 2021 · fda942f · fda942f
1 parent 67fa1f1
commit fda942f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
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,