BREAKING_CHANGE: deprecate Vertex SDK data science package

PiperOrigin-RevId: 640614179

Author: yinghsienwu

tests/unit/aiplatform/test_language_models.py

@@ -65,7 +65,6 @@
     prediction_service as gca_prediction_service_v1beta1,
 )
 
-import vertexai
 from vertexai.preview import (
     language_models as preview_language_models,
 )
@@ -4736,93 +4735,6 @@ def test_batch_prediction_for_text_embedding(self):
                 model_parameters={},
             )
 
-    def test_text_generation_top_level_from_pretrained_preview(self):
-        """Tests the text generation model."""
-        aiplatform.init(
-            project=_TEST_PROJECT,
-            location=_TEST_LOCATION,
-        )
-        with mock.patch.object(
-            target=model_garden_service_client.ModelGardenServiceClient,
-            attribute="get_publisher_model",
-            return_value=gca_publisher_model.PublisherModel(
-                _TEXT_BISON_PUBLISHER_MODEL_DICT
-            ),
-        ) as mock_get_publisher_model:
-            model = vertexai.preview.from_pretrained(
-                foundation_model_name="text-bison@001"
-            )
-
-            assert isinstance(model, preview_language_models.TextGenerationModel)
-
-        mock_get_publisher_model.assert_called_with(
-            name="publishers/google/models/text-bison@001", retry=base._DEFAULT_RETRY
-        )
-        assert mock_get_publisher_model.call_count == 1
-
-        assert (
-            model._model_resource_name
-            == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/text-bison@001"
-        )
-
-        # Test that methods on TextGenerationModel still work as expected
-        gca_predict_response = gca_prediction_service.PredictResponse()
-        gca_predict_response.predictions.append(_TEST_TEXT_GENERATION_PREDICTION)
-
-        with mock.patch.object(
-            target=prediction_service_client.PredictionServiceClient,
-            attribute="predict",
-            return_value=gca_predict_response,
-        ):
-            response = model.predict(
-                "What is the best recipe for banana bread? Recipe:",
-                max_output_tokens=128,
-                temperature=0.0,
-                top_p=1.0,
-                top_k=5,
-            )
-
-        assert response.text == _TEST_TEXT_GENERATION_PREDICTION["content"]
-        assert (
-            response.raw_prediction_response.predictions[0]
-            == _TEST_TEXT_GENERATION_PREDICTION
-        )
-        assert (
-            response.safety_attributes["Violent"]
-            == _TEST_TEXT_GENERATION_PREDICTION["safetyAttributes"]["scores"][0]
-        )
-
-    def test_text_embedding_top_level_from_pretrained_preview(self):
-        """Tests the text embedding model."""
-        aiplatform.init(
-            project=_TEST_PROJECT,
-            location=_TEST_LOCATION,
-        )
-        with mock.patch.object(
-            target=model_garden_service_client.ModelGardenServiceClient,
-            attribute="get_publisher_model",
-            return_value=gca_publisher_model.PublisherModel(
-                _TEXT_EMBEDDING_GECKO_PUBLISHER_MODEL_DICT
-            ),
-        ) as mock_get_publisher_model:
-            model = vertexai.preview.from_pretrained(
-                foundation_model_name="textembedding-gecko@001"
-            )
-
-            assert isinstance(model, preview_language_models.TextEmbeddingModel)
-
-            assert (
-                model._endpoint_name
-                == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/textembedding-gecko@001"
-            )
-
-        mock_get_publisher_model.assert_called_with(
-            name="publishers/google/models/textembedding-gecko@001",
-            retry=base._DEFAULT_RETRY,
-        )
-
-        assert mock_get_publisher_model.call_count == 1
-
 
 # TODO (b/285946649): add more test coverage before public preview release
 @pytest.mark.usefixtures("google_auth_mock")

tests/unit/aiplatform/test_vision_models.py

@@ -39,7 +39,7 @@
 from google.cloud.aiplatform.compat.types import (
     publisher_model as gca_publisher_model,
 )
-import vertexai
+
 from vertexai import vision_models as ga_vision_models
 from vertexai.preview import (
     vision_models as preview_vision_models,
@@ -221,48 +221,13 @@ def _get_image_generation_model(
 
         return model
 
-    def _get_preview_image_generation_model_top_level_from_pretrained(
-        self,
-    ) -> preview_vision_models.ImageGenerationModel:
-        """Gets the image generation model from the top-level vertexai.preview.from_pretrained method."""
-        aiplatform.init(
-            project=_TEST_PROJECT,
-            location=_TEST_LOCATION,
-        )
-        with mock.patch.object(
-            target=model_garden_service_client.ModelGardenServiceClient,
-            attribute="get_publisher_model",
-            return_value=gca_publisher_model.PublisherModel(
-                _IMAGE_GENERATION_PUBLISHER_MODEL_DICT
-            ),
-        ) as mock_get_publisher_model:
-            model = vertexai.preview.from_pretrained(
-                foundation_model_name="imagegeneration@002"
-            )
-
-        mock_get_publisher_model.assert_called_with(
-            name="publishers/google/models/imagegeneration@002",
-            retry=base._DEFAULT_RETRY,
-        )
-
-        assert mock_get_publisher_model.call_count == 1
-
-        return model
-
     def test_from_pretrained(self):
         model = self._get_image_generation_model()
         assert (
             model._endpoint_name
             == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/imagegeneration@002"
         )
 
-    def test_top_level_from_pretrained_preview(self):
-        model = self._get_preview_image_generation_model_top_level_from_pretrained()
-        assert (
-            model._endpoint_name
-            == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/imagegeneration@002"
-        )
-
     def test_generate_images(self):
         """Tests the image generation model."""
         model = self._get_image_generation_model()