samples: add model type for create_training_pipeline_image_classifica…

…tion_sample (#891)
googleapis · Dec 8, 2021 · a06da6d · a06da6d
1 parent 37ee0a1
commit a06da6d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 0 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,
+    model_type: str = "CLOUD",
     multi_label: bool = False,
     training_fraction_split: float = 0.8,
     validation_fraction_split: float = 0.1,
@@ -36,6 +37,7 @@ def create_training_pipeline_image_classification_sample(
 
     job = aiplatform.AutoMLImageTrainingJob(
         display_name=display_name,
+        model_type=model_type,
         prediction_type='classification',
         multi_label=multi_label
     )

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
@@ -45,6 +45,7 @@ def test_create_training_pipeline_image_classification_sample(
     )
     mock_get_automl_image_training_job.assert_called_once_with(
         display_name=constants.DISPLAY_NAME,
+        model_type=constants.MODEL_TYPE,
         multi_label=False,
         prediction_type='classification'
     )

diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
@@ -196,3 +196,4 @@
     "gs://bucket3/custom-training-python-package/my_app/trainer-0.1.tar.gz"
 )
 PYTHON_MODULE_NAME = "trainer.task"
+MODEL_TYPE = "CLOUD"