feat: support 32k text-generation and multilingual embedding models (#161)

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

Author: ashleyxuu

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,29 +25,37 @@
 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"
 
 
 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",
         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

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

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