diff --git a/bigframes/exceptions.py b/bigframes/exceptions.py index 00abb887b0..462bdf2bdd 100644 --- a/bigframes/exceptions.py +++ b/bigframes/exceptions.py @@ -69,3 +69,7 @@ class AmbiguousWindowWarning(Warning): class UnknownDataTypeWarning(Warning): """Data type is unknown.""" + + +class ApiDeprecationWarning(FutureWarning): + """The API has been deprecated.""" diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 53a9d40c6e..5fd713e465 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -21,9 +21,10 @@ import bigframes_vendored.constants as constants from google.cloud import bigquery +import typing_extensions import bigframes -from bigframes import clients +from bigframes import clients, exceptions from bigframes.core import blocks, log_adapter from bigframes.ml import base, core, globals, utils import bigframes.pandas as bpd @@ -407,12 +408,16 @@ def to_gbq(self, model_name: str, replace: bool = False) -> PaLM2TextGenerator: return new_model.session.read_gbq_model(model_name) +@typing_extensions.deprecated( + "PaLM2TextEmbeddingGenerator has been deprecated. Use TextEmbeddingGenerator(https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.TextEmbeddingGenerator) instead. ", + category=exceptions.ApiDeprecationWarning, +) @log_adapter.class_logger class PaLM2TextEmbeddingGenerator(base.BaseEstimator): """PaLM2 text embedding generator LLM model. .. note:: - Models in this class are outdated and going to be deprecated. To use the most updated text embedding models, go to the TextEmbeddingGenerator class. + PaLM2TextEmbeddingGenerator has been deprecated. Use TextEmbeddingGenerator(https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.TextEmbeddingGenerator) instead. Args: diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 3093a36534..f05efecfdd 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -14,6 +14,7 @@ import pytest +from bigframes import exceptions from bigframes.ml import llm from tests.system import utils @@ -383,3 +384,11 @@ def test_llm_gemini_pro_score_params(llm_fine_tune_df_default_index): ], index=6, ) + + +def test_palm2_text_embedding_deprecated(): + with pytest.warns(exceptions.ApiDeprecationWarning): + try: + llm.PaLM2TextEmbeddingGenerator() + except (Exception): + pass