feat: Add the recently added new model type "cloud_1" to the "AutoMLI…

…mageTrainingJob" in SDK. PiperOrigin-RevId: 512114110
googleapis · Feb 24, 2023 · 581939b · 581939b
1 parent 04ebc01
commit 581939b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py
@@ -81,8 +81,8 @@
 # NOTE: When adding a new prediction_type's, ensure it fits the pattern
 #       "automl_image_{prediction_type}_*" used by the YAML schemas on GCS
 AUTOML_IMAGE_PREDICTION_MODEL_TYPES = {
-    "classification": {"CLOUD"} | MOBILE_TF_MODEL_TYPES,
-    "object_detection": {"CLOUD_HIGH_ACCURACY_1", "CLOUD_LOW_LATENCY_1"}
+    "classification": {"CLOUD", "CLOUD_1"} | MOBILE_TF_MODEL_TYPES,
+    "object_detection": {"CLOUD_1", "CLOUD_HIGH_ACCURACY_1", "CLOUD_LOW_LATENCY_1"}
     | MOBILE_TF_MODEL_TYPES,
 }
 

diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py
@@ -5276,6 +5276,9 @@ def __init__(
                     "CLOUD" - Default for Image Classification.
                         A Model best tailored to be used within Google Cloud, and
                         which cannot be exported.
+                    "CLOUD_1" - A model type best tailored to be used within Google Cloud,
+                        which cannot be exported externally. Compared to the CLOUD model
+                        above, it is expected to have higher prediction accuracy.
                     "CLOUD_HIGH_ACCURACY_1" - Default for Image Object Detection.
                         A model best tailored to be used within Google Cloud, and
                         which cannot be exported. Expected to have a higher latency,