feat: add bigquery.ai.generate_text function (#2433)

* Added the API for ai.generate_text
* Fixed SQL syntax bug of ai.generate_embedding
* Refactored the code base to keep util functions organized.
Fixes b/481092205

Author: sycai

bigframes/bigquery/_operations/ai.py

@@ -26,6 +26,7 @@
 frombigframesimportclients, dataframe, dtypes
 frombigframesimportpandasasbpd
 frombigframesimportseries, session
+frombigframes.bigquery._operationsimportutilsasbq_utils
 frombigframes.coreimportconvert
 frombigframes.core.loggingimportlog_adapter
 importbigframes.core.sql.literals
@@ -391,7 +392,7 @@ def generate_double(
 
 @log_adapter.method_logger(custom_base_name="bigquery_ai")
 defgenerate_embedding(
-model_name: str,
+model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
 data: Union[dataframe.DataFrame, series.Series, pd.DataFrame, pd.Series],
 *,
 output_dimensionality: Optional[int] =None,
@@ -415,9 +416,8 @@ def generate_embedding(
  ... ) # doctest: +SKIP
 
  Args:
- model_name (str):
- The name of a remote model from Vertex AI, such as the
- multimodalembedding@001 model.
+ model (bigframes.ml.base.BaseEstimator or str):
+ The model to use for text embedding.
  data (bigframes.pandas.DataFrame or bigframes.pandas.Series):
  The data to generate embeddings for. If a Series is provided, it is
  treated as the 'content' column. If a DataFrame is provided, it
@@ -454,20 +454,9 @@ def generate_embedding(
  <https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-generate-embedding#output>`_
  for details.
  """
-ifisinstance(data, (pd.DataFrame, pd.Series)):
-data=bpd.read_pandas(data)
-
-ifisinstance(data, series.Series):
-data=data.copy()
-data.name="content"
-data_df=data.to_frame()
-elifisinstance(data, dataframe.DataFrame):
-data_df=data
-else:
-raiseValueError(f"Unsupported data type: {type(data)}")
-
-# We need to get the SQL for the input data to pass as a subquery to the TVF
-source_sql=data_df.sql
+data=_to_dataframe(data, series_rename="content")
+model_name, session=bq_utils.get_model_name_and_session(model, data)
+table_sql=bq_utils.to_sql(data)
 
 struct_fields: Dict[str, bigframes.core.sql.literals.STRUCT_VALUES] = {}
 ifoutput_dimensionalityisnotNone:
@@ -488,12 +477,128 @@ def generate_embedding(
  SELECT *
  FROM AI.GENERATE_EMBEDDING(
  MODEL `{model_name}`,
- ({source_sql}),
-{bigframes.core.sql.literals.struct_literal(struct_fields)})
+ ({table_sql}),
+{bigframes.core.sql.literals.struct_literal(struct_fields)}
  )
  """
 
-returndata_df._session.read_gbq(query)
+ifsessionisNone:
+returnbpd.read_gbq_query(query)
+else:
+returnsession.read_gbq_query(query)
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_ai")
+defgenerate_text(
+model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
+data: Union[dataframe.DataFrame, series.Series, pd.DataFrame, pd.Series],
+*,
+temperature: Optional[float] =None,
+max_output_tokens: Optional[int] =None,
+top_k: Optional[int] =None,
+top_p: Optional[float] =None,
+stop_sequences: Optional[List[str]] =None,
+ground_with_google_search: Optional[bool] =None,
+request_type: Optional[str] =None,
+) ->dataframe.DataFrame:
+"""
+ Generates text using a BigQuery ML model.
+
+ See the `BigQuery ML GENERATE_TEXT function syntax
+ <https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-text>`_
+ for additional reference.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> import bigframes.bigquery as bbq
+ >>> df = bpd.DataFrame({"prompt": ["write a poem about apples"]})
+ >>> bbq.ai.generate_text(
+ ... "project.dataset.model_name",
+ ... df
+ ... ) # doctest: +SKIP
+
+ Args:
+ model (bigframes.ml.base.BaseEstimator or str):
+ The model to use for text generation.
+ data (bigframes.pandas.DataFrame or bigframes.pandas.Series):
+ The data to generate embeddings for. If a Series is provided, it is
+ treated as the 'content' column. If a DataFrame is provided, it
+ must contain a 'content' column, or you must rename the column you
+ wish to embed to 'content'.
+ temperature (float, optional):
+ A FLOAT64 value that is used for sampling promiscuity. The value
+ must be in the range ``[0.0, 1.0]``. A lower temperature works well
+ for prompts that expect a more deterministic and less open-ended
+ or creative response, while a higher temperature can lead to more
+ diverse or creative results. A temperature of ``0`` is
+ deterministic, meaning that the highest probability response is
+ always selected.
+ max_output_tokens (int, optional):
+ An INT64 value that sets the maximum number of tokens in the
+ generated text.
+ top_k (int, optional):
+ An INT64 value that changes how the model selects tokens for
+ output. A ``top_k`` of ``1`` means the next selected token is the
+ most probable among all tokens in the model's vocabulary. A
+ ``top_k`` of ``3`` means that the next token is selected from
+ among the three most probable tokens by using temperature. The
+ default value is ``40``.
+ top_p (float, optional):
+ A FLOAT64 value that changes how the model selects tokens for
+ output. Tokens are selected from most probable to least probable
+ until the sum of their probabilities equals the ``top_p`` value.
+ For example, if tokens A, B, and C have a probability of 0.3, 0.2,
+ and 0.1 and the ``top_p`` value is ``0.5``, then the model will
+ select either A or B as the next token by using temperature. The
+ default value is ``0.95``.
+ stop_sequences (List[str], optional):
+ An ARRAY<STRING> value that contains the stop sequences for the model.
+ ground_with_google_search (bool, optional):
+ A BOOL value that determines whether to ground the model with Google Search.
+ request_type (str, optional):
+ A STRING value that contains the request type for the model.
+
+ Returns:
+ bigframes.pandas.DataFrame:
+ The generated text.
+ """
+data=_to_dataframe(data, series_rename="prompt")
+model_name, session=bq_utils.get_model_name_and_session(model, data)
+table_sql=bq_utils.to_sql(data)
+
+struct_fields: Dict[
+str,
+Union[str, int, float, bool, Mapping[str, str], List[str], Mapping[str, Any]],
+ ] = {}
+iftemperatureisnotNone:
+struct_fields["TEMPERATURE"] =temperature
+ifmax_output_tokensisnotNone:
+struct_fields["MAX_OUTPUT_TOKENS"] =max_output_tokens
+iftop_kisnotNone:
+struct_fields["TOP_K"] =top_k
+iftop_pisnotNone:
+struct_fields["TOP_P"] =top_p
+ifstop_sequencesisnotNone:
+struct_fields["STEP_SEQUENCES"] =stop_sequences
+ifground_with_google_searchisnotNone:
+struct_fields["GROUND_WITH_GOOGLE_SEARCH"] =ground_with_google_search
+ifrequest_typeisnotNone:
+struct_fields["REQUEST_TYPE"] =request_type
+
+query=f"""
+ SELECT *
+ FROM AI.GENERATE_TEXT(
+ MODEL `{model_name}`,
+ ({table_sql}),
+{bigframes.core.sql.literals.struct_literal(struct_fields)}
+ )
+ """
+
+ifsessionisNone:
+returnbpd.read_gbq_query(query)
+else:
+returnsession.read_gbq_query(query)
 
 
 @log_adapter.method_logger(custom_base_name="bigquery_ai")
@@ -811,3 +916,20 @@ def _resolve_connection_id(series: series.Series, connection_id: str | None):
 series._session._project,
 series._session._location,
  )
+
+
+def_to_dataframe(
+data: Union[dataframe.DataFrame, series.Series, pd.DataFrame, pd.Series],
+series_rename: str,
+) ->dataframe.DataFrame:
+ifisinstance(data, (pd.DataFrame, pd.Series)):
+data=bpd.read_pandas(data)
+
+ifisinstance(data, series.Series):
+data=data.copy()
+data.name=series_rename
+returndata.to_frame()
+elifisinstance(data, dataframe.DataFrame):
+returndata
+
+raiseValueError(f"Unsupported data type: {type(data)}")

bigframes/bigquery/_operations/ml.py

@@ -14,66 +14,20 @@
 
 from __future__ importannotations
 
-fromtypingimportcast, List, Mapping, Optional, Union
+fromtypingimportList, Mapping, Optional, Union
 
 importbigframes_vendored.constants
 importgoogle.cloud.bigquery
 importpandasaspd
 
+frombigframes.bigquery._operationsimportutils
 importbigframes.core.logging.log_adapteraslog_adapter
 importbigframes.core.sql.ml
 importbigframes.dataframeasdataframe
 importbigframes.ml.base
 importbigframes.session
 
 
-# Helper to convert DataFrame to SQL string
-def_to_sql(df_or_sql: Union[pd.DataFrame, dataframe.DataFrame, str]) ->str:
-importbigframes.pandasasbpd
-
-ifisinstance(df_or_sql, str):
-returndf_or_sql
-
-ifisinstance(df_or_sql, pd.DataFrame):
-bf_df=bpd.read_pandas(df_or_sql)
-else:
-bf_df=cast(dataframe.DataFrame, df_or_sql)
-
-# Cache dataframes to make sure base table is not a snapshot.
-# Cached dataframe creates a full copy, never uses snapshot.
-# This is a workaround for internal issue b/310266666.
-bf_df.cache()
-sql, _, _=bf_df._to_sql_query(include_index=False)
-returnsql
-
-
-def_get_model_name_and_session(
-model: Union[bigframes.ml.base.BaseEstimator, str, pd.Series],
-# Other dataframe arguments to extract session from
-*dataframes: Optional[Union[pd.DataFrame, dataframe.DataFrame, str]],
-) ->tuple[str, Optional[bigframes.session.Session]]:
-ifisinstance(model, pd.Series):
-try:
-model_ref=model["modelReference"]
-model_name=f"{model_ref['projectId']}.{model_ref['datasetId']}.{model_ref['modelId']}"# type: ignore
-exceptKeyError:
-raiseValueError("modelReference must be present in the pandas Series.")
-elifisinstance(model, str):
-model_name=model
-else:
-ifmodel._bqml_modelisNone:
-raiseValueError("Model must be fitted to be used in ML operations.")
-returnmodel._bqml_model.model_name, model._bqml_model.session
-
-session=None
-fordfindataframes:
-ifisinstance(df, dataframe.DataFrame):
-session=df._session
-break
-
-returnmodel_name, session
-
-
 def_get_model_metadata(
 *,
 bqclient: google.cloud.bigquery.Client,
@@ -143,8 +97,12 @@ def create_model(
  """
 importbigframes.pandasasbpd
 
-training_data_sql=_to_sql(training_data) iftraining_dataisnotNoneelseNone
-custom_holiday_sql=_to_sql(custom_holiday) ifcustom_holidayisnotNoneelseNone
+training_data_sql= (
+utils.to_sql(training_data) iftraining_dataisnotNoneelseNone
+ )
+custom_holiday_sql= (
+utils.to_sql(custom_holiday) ifcustom_holidayisnotNoneelseNone
+ )
 
 # Determine session from DataFrames if not provided
 ifsessionisNone:
@@ -227,8 +185,8 @@ def evaluate(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model, input_)
-table_sql=_to_sql(input_) ifinput_isnotNoneelseNone
+model_name, session=utils.get_model_name_and_session(model, input_)
+table_sql=utils.to_sql(input_) ifinput_isnotNoneelseNone
 
 sql=bigframes.core.sql.ml.evaluate(
 model_name=model_name,
@@ -281,8 +239,8 @@ def predict(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model, input_)
-table_sql=_to_sql(input_)
+model_name, session=utils.get_model_name_and_session(model, input_)
+table_sql=utils.to_sql(input_)
 
 sql=bigframes.core.sql.ml.predict(
 model_name=model_name,
@@ -340,8 +298,8 @@ def explain_predict(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model, input_)
-table_sql=_to_sql(input_)
+model_name, session=utils.get_model_name_and_session(model, input_)
+table_sql=utils.to_sql(input_)
 
 sql=bigframes.core.sql.ml.explain_predict(
 model_name=model_name,
@@ -383,7 +341,7 @@ def global_explain(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model)
+model_name, session=utils.get_model_name_and_session(model)
 sql=bigframes.core.sql.ml.global_explain(
 model_name=model_name,
 class_level_explain=class_level_explain,
@@ -419,8 +377,8 @@ def transform(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model, input_)
-table_sql=_to_sql(input_)
+model_name, session=utils.get_model_name_and_session(model, input_)
+table_sql=utils.to_sql(input_)
 
 sql=bigframes.core.sql.ml.transform(
 model_name=model_name,
@@ -500,8 +458,8 @@ def generate_text(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model, input_)
-table_sql=_to_sql(input_)
+model_name, session=utils.get_model_name_and_session(model, input_)
+table_sql=utils.to_sql(input_)
 
 sql=bigframes.core.sql.ml.generate_text(
 model_name=model_name,
@@ -565,8 +523,8 @@ def generate_embedding(
  """
 importbigframes.pandasasbpd
 
-model_name, session=_get_model_name_and_session(model, input_)
-table_sql=_to_sql(input_)
+model_name, session=utils.get_model_name_and_session(model, input_)
+table_sql=utils.to_sql(input_)
 
 sql=bigframes.core.sql.ml.generate_embedding(
 model_name=model_name,