Fixing vertexai imports.

PiperOrigin-RevId: 633213566 Change-Id: I2a90caeecd787f19f4e1875d23a5fcb3d9c39f86
google-deepmind · May 13, 2024 · 43c3b67 · 43c3b67
1 parent 7f2ef88
commit 43c3b67
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/onetwo/backends/vertexai_api.py b/onetwo/backends/vertexai_api.py
@@ -24,9 +24,9 @@
 from typing import Any, Final, Iterable, cast
 
 from absl import logging
-from google.cloud.aiplatform import vertexai
-from google.cloud.aiplatform.vertexai import generative_models
-from google.cloud.aiplatform.vertexai import language_models
+import vertexai
+from vertexai import generative_models
+from vertexai import language_models
 import immutabledict
 from onetwo.backends import backends_base
 from onetwo.builtins import formatting

diff --git a/onetwo/backends/vertexai_api_test.py b/onetwo/backends/vertexai_api_test.py
@@ -21,12 +21,12 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from google.cloud.aiplatform import vertexai
-from google.cloud.aiplatform.aiplatform_v1beta1.types import (
+import vertexai
+from google.cloud.aiplatform_v1beta1.types import (
     prediction_service as gapic_prediction_service_types,
 )
-from google.cloud.aiplatform.vertexai import generative_models
-from google.cloud.aiplatform.vertexai import language_models
+from vertexai import generative_models
+from vertexai import language_models
 from onetwo.backends import vertexai_api
 from onetwo.builtins import llm
 from onetwo.core import caching