From 43c3b67e4b479e49e31baefb304d1487b4c893bb Mon Sep 17 00:00:00 2001 From: Ilya Tolstikhin Date: Mon, 13 May 2024 08:07:05 -0700 Subject: [PATCH] Fixing vertexai imports. PiperOrigin-RevId: 633213566 Change-Id: I2a90caeecd787f19f4e1875d23a5fcb3d9c39f86 --- onetwo/backends/vertexai_api.py | 6 +++--- onetwo/backends/vertexai_api_test.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/onetwo/backends/vertexai_api.py b/onetwo/backends/vertexai_api.py index 8cdb831..1191975 100644 --- 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 index 5fbb9e4..27e987d 100644 --- 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