From 5d0734792e98afd889d27e0b8a18ac60cb0fc50d Mon Sep 17 00:00:00 2001 From: Ashley Xu Date: Tue, 31 Oct 2023 21:23:25 +0000 Subject: [PATCH 1/4] feat: support 32k text-generation and embedding models --- bigframes/ml/llm.py | 36 +++++++++++++------ tests/system/small/ml/conftest.py | 18 ++++++++++ tests/system/small/ml/test_llm.py | 57 +++++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+), 10 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index d78f467537..60f64dd358 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_32K_MODEL_CODE = "textembedding-gecko-multilingual" _EMBED_TEXT_RESULT_COLUMN = "text_embedding" @@ -45,9 +47,11 @@ class PaLM2TextGenerator(base.Predictor): def __init__( self, + model_name: Literal["text-bison", "text-bison-32k"] = "text-bison-32k", 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,10 +84,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 +126,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]. @@ -184,9 +192,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 +231,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_32K_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..d467bb8c71 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_32k_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..e44d80ec02 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 @@ -139,6 +169,20 @@ def test_create_text_embedding_generator_model_defaults(bq_connection): assert model._bqml_model is not None +def test_create_text_embedding_generator_32k_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 +196,19 @@ def test_embedding_generator_predict_success( assert value.size == 768 +@pytest.mark.flaky(retries=2, delay=120) +def test_embedding_generator_32k_predict_success( + palm2_embedding_generator_32k_model, llm_text_df +): + df = palm2_embedding_generator_32k_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 From ef0116a95a301f8b0328bfeb175eeb22cf463b9e Mon Sep 17 00:00:00 2001 From: Ashley Xu Date: Tue, 31 Oct 2023 21:23:25 +0000 Subject: [PATCH 2/4] feat: support 32k text-generation and embedding multilingual models --- bigframes/ml/llm.py | 30 ++++++++++++++++++-- tests/system/small/ml/conftest.py | 9 +++++- tests/system/small/ml/test_llm.py | 46 ++++++++++++++++++++++++++++--- 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 60f64dd358..e24121dc70 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -26,10 +26,11 @@ _REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" _REMOTE_TEXT_GENERATOR_32K_MODEL_CODE = "text-bison-32k" +_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_32K_MODEL_CODE = "textembedding-gecko-multilingual" +_REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE = "textembedding-gecko-multilingual" _EMBED_TEXT_RESULT_COLUMN = "text_embedding" @@ -37,10 +38,14 @@ 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. """ @@ -48,9 +53,11 @@ class PaLM2TextGenerator(base.Predictor): def __init__( self, model_name: Literal["text-bison", "text-bison-32k"] = "text-bison-32k", + model_name: Literal["text-bison", "text-bison-32k"] = "text-bison-32k", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): + self.model_name = model_name self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( @@ -92,6 +99,14 @@ def _create_bqml_model(self): options = { "endpoint": _REMOTE_TEXT_GENERATOR_32K_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 @@ -127,6 +142,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]. @@ -183,6 +199,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): @@ -195,9 +215,13 @@ def __init__( model_name: Literal[ "textembedding-gecko", "textembedding-gecko-multilingual" ] = "textembedding-gecko", + 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.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( @@ -237,7 +261,7 @@ def _create_bqml_model(self): } else: options = { - "endpoint": _REMOTE_EMBEDDING_GENERATOR_32K_MODEL_CODE, + "endpoint": _REMOTE_EMBEDDING_GENERATOR_MUlTILINGUAL_MODEL_CODE, } return self._bqml_model_factory.create_remote_model( diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index d467bb8c71..76ea6c4c08 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -220,6 +220,13 @@ def palm2_text_generator_32k_model(session, bq_connection) -> llm.PaLM2TextGener ) +@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 @@ -237,7 +244,7 @@ def palm2_embedding_generator_model( @pytest.fixture(scope="session") -def palm2_embedding_generator_32k_model( +def palm2_embedding_generator_multilingual_model( session, bq_connection ) -> llm.PaLM2TextEmbeddingGenerator: return llm.PaLM2TextEmbeddingGenerator( diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index e44d80ec02..2b27b7a9bb 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -32,6 +32,12 @@ def test_create_text_generator_32k_model(palm2_text_generator_32k_model): assert palm2_text_generator_32k_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 @@ -78,6 +84,30 @@ def test_create_text_generator_32k_model_default_session( 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 @@ -157,6 +187,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 @@ -169,7 +207,7 @@ def test_create_text_embedding_generator_model_defaults(bq_connection): assert model._bqml_model is not None -def test_create_text_embedding_generator_32k_model_defaults(bq_connection): +def test_create_text_embedding_generator_multilingual_model_defaults(bq_connection): import bigframes.pandas as bpd bpd.close_session() @@ -197,10 +235,10 @@ def test_embedding_generator_predict_success( @pytest.mark.flaky(retries=2, delay=120) -def test_embedding_generator_32k_predict_success( - palm2_embedding_generator_32k_model, llm_text_df +def test_embedding_generator_multilingual_predict_success( + palm2_embedding_generator_multilingual_model, llm_text_df ): - df = palm2_embedding_generator_32k_model.predict(llm_text_df).to_pandas() + 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"] From 1011fb894c572ed81c3569b6ff0184bae5a92482 Mon Sep 17 00:00:00 2001 From: Ashley Xu Date: Tue, 31 Oct 2023 22:14:11 +0000 Subject: [PATCH 3/4] fix: fix the merging conflicts --- bigframes/ml/llm.py | 15 --------------- tests/system/small/ml/conftest.py | 7 ------- tests/system/small/ml/test_llm.py | 30 ------------------------------ 3 files changed, 52 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index e24121dc70..53168b29b5 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -53,11 +53,9 @@ class PaLM2TextGenerator(base.Predictor): def __init__( self, model_name: Literal["text-bison", "text-bison-32k"] = "text-bison-32k", - model_name: Literal["text-bison", "text-bison-32k"] = "text-bison-32k", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): - self.model_name = model_name self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( @@ -99,15 +97,6 @@ def _create_bqml_model(self): options = { "endpoint": _REMOTE_TEXT_GENERATOR_32K_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 ) @@ -215,13 +204,9 @@ def __init__( model_name: Literal[ "textembedding-gecko", "textembedding-gecko-multilingual" ] = "textembedding-gecko", - 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.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = clients.BqConnectionManager( diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 76ea6c4c08..c11445b79a 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -220,13 +220,6 @@ def palm2_text_generator_32k_model(session, bq_connection) -> llm.PaLM2TextGener ) -@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 diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 2b27b7a9bb..79d3c40317 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -32,12 +32,6 @@ def test_create_text_generator_32k_model(palm2_text_generator_32k_model): assert palm2_text_generator_32k_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 @@ -84,30 +78,6 @@ def test_create_text_generator_32k_model_default_session( 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 From f782f4fc9daa4cce6b07ea577b05b3a50f2a0cab Mon Sep 17 00:00:00 2001 From: Ashley Xu Date: Wed, 1 Nov 2023 21:04:52 +0000 Subject: [PATCH 4/4] fix: address the comments --- bigframes/ml/llm.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 53168b29b5..2e5a9a1e5e 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -26,7 +26,6 @@ _REMOTE_TEXT_GENERATOR_MODEL_CODE = "CLOUD_AI_LARGE_LANGUAGE_MODEL_V1" _REMOTE_TEXT_GENERATOR_32K_MODEL_CODE = "text-bison-32k" -_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" @@ -52,7 +51,7 @@ class PaLM2TextGenerator(base.Predictor): def __init__( self, - model_name: Literal["text-bison", "text-bison-32k"] = "text-bison-32k", + model_name: Literal["text-bison", "text-bison-32k"] = "text-bison", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): @@ -131,7 +130,6 @@ 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].