diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index d78f467537..2e5a9a1e5e 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import cast, Optional, Union +from typing import cast, Literal, Optional, Union import bigframes from bigframes import clients, constants @@ -25,9 +25,11 @@ import bigframes.pandas as bpd _REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" +_REMOTE_TEXT_GENERATOR_32K_MODEL_CODE = "text-bison-32k" _TEXT_GENERATE_RESULT_COLUMN = "ml_generate_text_llm_result" _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE = "CLOUD_AI_TEXT_EMBEDDING_MODEL_V1" +_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE = "textembedding-gecko-multilingual" _EMBED_TEXT_RESULT_COLUMN = "text_embedding" @@ -35,19 +37,25 @@ class PaLM2TextGenerator(base.Predictor): """PaLM2 text generator LLM model. Args: + model_name (str, Default to "text-bison"): + The model for natural language tasks. “text-bison” returns model fine-tuned to follow natural language instructions + and is suitable for a variety of language tasks. "text-bison-32k" supports up to 32k tokens per request. + Default to "text-bison". session (bigframes.Session or None): BQ session to create the model. If None, use the global default session. connection_name (str or None): - connection to connect with remote service. str of the format ... + Connection to connect with remote service. str of the format ... if None, use default connection in session context. BigQuery DataFrame will try to create the connection and attach permission if the connection isn't fully setup. """ def __init__( self, + model_name: Literal["text-bison", "text-bison-32k"] = "text-bison", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): + self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient @@ -80,11 +88,14 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - - options = { - "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, - } - + if self.model_name == "text-bison": + options = { + "remote_service_type": _REMOTE_TEXT_GENERATOR_MODEL_CODE, + } + else: + options = { + "endpoint": _REMOTE_TEXT_GENERATOR_32K_MODEL_CODE, + } return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options ) @@ -118,7 +129,7 @@ def predict( top_k (int, default 40): Top-k changes how the model selects tokens for output. A top-k of 1 means the selected token is the most probable among all tokens - in the model’s vocabulary (also called greedy decoding), while a top-k of 3 means that the next token is selected from among the 3 most probable tokens (using temperature). + in the model's vocabulary (also called greedy decoding), while a top-k of 3 means that the next token is selected from among the 3 most probable tokens (using temperature). For each token selection step, the top K tokens with the highest probabilities are sampled. Then tokens are further filtered based on topP with the final token selected using temperature sampling. Specify a lower value for less random responses and a higher value for more random responses. Default 40. Possible values [1, 40]. @@ -175,6 +186,10 @@ class PaLM2TextEmbeddingGenerator(base.Predictor): """PaLM2 text embedding generator LLM model. Args: + model_name (str, Default to "textembedding-gecko"): + The model for text embedding. “textembedding-gecko” returns model embeddings for text inputs. + "textembedding-gecko-multilingual" returns model embeddings for text inputs which support over 100 languages + Default to "textembedding-gecko". session (bigframes.Session or None): BQ session to create the model. If None, use the global default session. connection_name (str or None): @@ -184,9 +199,13 @@ class PaLM2TextEmbeddingGenerator(base.Predictor): def __init__( self, + model_name: Literal[ + "textembedding-gecko", "textembedding-gecko-multilingual" + ] = "textembedding-gecko", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): + self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( self.session.bqconnectionclient, self.session.resourcemanagerclient @@ -219,10 +238,14 @@ def _create_bqml_model(self): connection_id=connection_name_parts[2], iam_role="aiplatform.user", ) - - options = { - "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, - } + if self.model_name == "textembedding-gecko": + options = { + "remote_service_type": _REMOTE_EMBEDDING_GENERATOR_MODEL_CODE, + } + else: + options = { + "endpoint": _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE, + } return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 1dd1c813b8..c11445b79a 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -213,6 +213,13 @@ def palm2_text_generator_model(session, bq_connection) -> llm.PaLM2TextGenerator return llm.PaLM2TextGenerator(session=session, connection_name=bq_connection) +@pytest.fixture(scope="session") +def palm2_text_generator_32k_model(session, bq_connection) -> llm.PaLM2TextGenerator: + return llm.PaLM2TextGenerator( + model_name="text-bison-32k", session=session, connection_name=bq_connection + ) + + @pytest.fixture(scope="function") def ephemera_palm2_text_generator_model( session, bq_connection @@ -229,6 +236,17 @@ def palm2_embedding_generator_model( ) +@pytest.fixture(scope="session") +def palm2_embedding_generator_multilingual_model( + session, bq_connection +) -> llm.PaLM2TextEmbeddingGenerator: + return llm.PaLM2TextEmbeddingGenerator( + model_name="textembedding-gecko-multilingual", + session=session, + connection_name=bq_connection, + ) + + @pytest.fixture(scope="session") def time_series_bqml_arima_plus_model( session, time_series_arima_plus_model_name diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index b7257dde1b..79d3c40317 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -26,6 +26,12 @@ def test_create_text_generator_model(palm2_text_generator_model): assert palm2_text_generator_model._bqml_model is not None +def test_create_text_generator_32k_model(palm2_text_generator_32k_model): + # Model creation doesn't return error + assert palm2_text_generator_32k_model is not None + assert palm2_text_generator_32k_model._bqml_model is not None + + @pytest.mark.flaky(retries=2, delay=120) def test_create_text_generator_model_default_session(bq_connection, llm_text_pandas_df): import bigframes.pandas as bpd @@ -48,6 +54,30 @@ def test_create_text_generator_model_default_session(bq_connection, llm_text_pan assert all(series.str.len() > 20) +@pytest.mark.flaky(retries=2, delay=120) +def test_create_text_generator_32k_model_default_session( + bq_connection, llm_text_pandas_df +): + import bigframes.pandas as bpd + + bpd.close_session() + bpd.options.bigquery.bq_connection = bq_connection + bpd.options.bigquery.location = "us" + + model = llm.PaLM2TextGenerator(model_name="text-bison-32k") + assert model is not None + assert model._bqml_model is not None + assert model.connection_name.casefold() == "bigframes-dev.us.bigframes-rf-conn" + + llm_text_df = bpd.read_pandas(llm_text_pandas_df) + + df = model.predict(llm_text_df).to_pandas() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "ml_generate_text_llm_result" in df.columns + series = df["ml_generate_text_llm_result"] + assert all(series.str.len() > 20) + + @pytest.mark.flaky(retries=2, delay=120) def test_create_text_generator_model_default_connection(llm_text_pandas_df): from bigframes import _config @@ -127,6 +157,14 @@ def test_create_embedding_generator_model(palm2_embedding_generator_model): assert palm2_embedding_generator_model._bqml_model is not None +def test_create_embedding_generator_multilingual_model( + palm2_embedding_generator_multilingual_model, +): + # Model creation doesn't return error + assert palm2_embedding_generator_multilingual_model is not None + assert palm2_embedding_generator_multilingual_model._bqml_model is not None + + def test_create_text_embedding_generator_model_defaults(bq_connection): import bigframes.pandas as bpd @@ -139,6 +177,20 @@ def test_create_text_embedding_generator_model_defaults(bq_connection): assert model._bqml_model is not None +def test_create_text_embedding_generator_multilingual_model_defaults(bq_connection): + import bigframes.pandas as bpd + + bpd.close_session() + bpd.options.bigquery.bq_connection = bq_connection + bpd.options.bigquery.location = "us" + + model = llm.PaLM2TextEmbeddingGenerator( + model_name="textembedding-gecko-multilingual" + ) + assert model is not None + assert model._bqml_model is not None + + @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_success( palm2_embedding_generator_model, llm_text_df @@ -152,6 +204,19 @@ def test_embedding_generator_predict_success( assert value.size == 768 +@pytest.mark.flaky(retries=2, delay=120) +def test_embedding_generator_multilingual_predict_success( + palm2_embedding_generator_multilingual_model, llm_text_df +): + df = palm2_embedding_generator_multilingual_model.predict(llm_text_df).to_pandas() + TestCase().assertSequenceEqual(df.shape, (3, 1)) + assert "text_embedding" in df.columns + series = df["text_embedding"] + value = series[0] + assert isinstance(value, np.ndarray) + assert value.size == 768 + + @pytest.mark.flaky(retries=2, delay=120) def test_embedding_generator_predict_series_success( palm2_embedding_generator_model, llm_text_df