feat: GenAI - Added support for "models/<model ID>" model name format

PiperOrigin-RevId: 592693755
googleapis · Dec 21, 2023 · ab21feb · ab21feb
1 parent c033c59
commit ab21feb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -152,6 +152,8 @@ def __init__(
         """
         if "/" not in model_name:
             model_name = "publishers/google/models/" + model_name
+        if model_name.startswith("models/"):
+            model_name = "publishers/google/" + model_name
 
         project = aiplatform_initializer.global_config.project
         location = aiplatform_initializer.global_config.location