feat: support 32k text-generation and embedding multilingual models

Author: ashleyxuu

bigframes/ml/llm.py

@@ -26,31 +26,38 @@
 
 _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"
 
 
 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 <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
+            Connection to connect with remote service. str of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
             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-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(

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(

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"]