remove params

Author: GarrettWu

bigframes/bigquery/_operations/ml.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Any, cast, List, Mapping, Optional, Union
+from typing import cast, List, Mapping, Optional, Union
 
 import bigframes_vendored.constants
 import google.cloud.bigquery
@@ -443,10 +443,8 @@ def generate_text(
     top_k: Optional[int] = None,
     top_p: Optional[float] = None,
     flatten_json_output: Optional[bool] = None,
-    safety_settings: Optional[Mapping[str, str]] = None,
     stop_sequences: Optional[List[str]] = None,
     ground_with_google_search: Optional[bool] = None,
-    model_params: Optional[Mapping[str, Any]] = None,
     request_type: Optional[str] = None,
 ) -> dataframe.DataFrame:
     """
@@ -489,16 +487,10 @@ def generate_text(
             default value is ``0.95``.
         flatten_json_output (bool, optional):
             A BOOL value that determines the content of the generated JSON column.
-        safety_settings (Mapping[str, str], optional):
-            A STRUCT value that contains the safety settings for the model.
-            The STRUCT must have a ``category`` field of type STRING and a
-            ``threshold`` field of type STRING.
         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.
-        model_params (Mapping[str, Any], optional):
-            A JSON value that contains the parameters for the model.
         request_type (str, optional):
             A STRING value that contains the request type for the model.
 
@@ -519,10 +511,8 @@ def generate_text(
         top_k=top_k,
         top_p=top_p,
         flatten_json_output=flatten_json_output,
-        safety_settings=safety_settings,
         stop_sequences=stop_sequences,
         ground_with_google_search=ground_with_google_search,
-        model_params=model_params,
         request_type=request_type,
     )
 

bigframes/core/sql/ml.py

@@ -264,10 +264,8 @@ def generate_text(
     top_k: Optional[int] = None,
     top_p: Optional[float] = None,
     flatten_json_output: Optional[bool] = None,
-    safety_settings: Optional[Mapping[str, str]] = None,
     stop_sequences: Optional[List[str]] = None,
     ground_with_google_search: Optional[bool] = None,
-    model_params: Optional[Mapping[str, Any]] = None,
     request_type: Optional[str] = None,
 ) -> str:
     """Encode the ML.GENERATE_TEXT statement.
@@ -287,14 +285,10 @@ def generate_text(
         struct_options["top_p"] = top_p
     if flatten_json_output is not None:
         struct_options["flatten_json_output"] = flatten_json_output
-    if safety_settings is not None:
-        struct_options["safety_settings"] = safety_settings
     if stop_sequences is not None:
         struct_options["stop_sequences"] = stop_sequences
     if ground_with_google_search is not None:
         struct_options["ground_with_google_search"] = ground_with_google_search
-    if model_params is not None:
-        struct_options["model_params"] = model_params
     if request_type is not None:
         struct_options["request_type"] = request_type
 

tests/unit/bigquery/test_ml.py

@@ -182,10 +182,8 @@ def test_generate_text_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mo
         top_k=20,
         top_p=0.9,
         flatten_json_output=True,
-        safety_settings={"hate_speech": "BLOCK_ONLY_HIGH"},
         stop_sequences=["a", "b"],
         ground_with_google_search=True,
-        model_params={"param1": "value1"},
         request_type="TYPE",
     )
     read_pandas_mock.assert_called_once()
@@ -199,10 +197,6 @@ def test_generate_text_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mo
     assert "20 AS top_k" in generated_sql
     assert "0.9 AS top_p" in generated_sql
     assert "true AS flatten_json_output" in generated_sql
-    assert (
-        "STRUCT('BLOCK_ONLY_HIGH' AS hate_speech) AS safety_settings" in generated_sql
-    )
     assert "['a', 'b'] AS stop_sequences" in generated_sql
     assert "true AS ground_with_google_search" in generated_sql
-    assert """JSON'{"param1": "value1"}' AS model_params""" in generated_sql
     assert "'TYPE' AS request_type" in generated_sql

tests/unit/core/sql/snapshots/test_ml/test_generate_text_model_with_options/generate_text_model_with_options.sql

@@ -1 +1 @@
-SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM new_data), STRUCT(0.5 AS temperature, 128 AS max_output_tokens, 20 AS top_k, 0.9 AS top_p, true AS flatten_json_output, STRUCT('BLOCK_ONLY_HIGH' AS hate_speech) AS safety_settings, ['a', 'b'] AS stop_sequences, true AS ground_with_google_search, JSON'{"param1": "value1"}' AS model_params, 'TYPE' AS request_type))
+SELECT * FROM ML.GENERATE_TEXT(MODEL `my_project.my_dataset.my_model`, (SELECT * FROM new_data), STRUCT(0.5 AS temperature, 128 AS max_output_tokens, 20 AS top_k, 0.9 AS top_p, true AS flatten_json_output, ['a', 'b'] AS stop_sequences, true AS ground_with_google_search, 'TYPE' AS request_type))

tests/unit/core/sql/test_ml.py

@@ -196,10 +196,8 @@ def test_generate_text_model_with_options(snapshot):
         top_k=20,
         top_p=0.9,
         flatten_json_output=True,
-        safety_settings={"hate_speech": "BLOCK_ONLY_HIGH"},
         stop_sequences=["a", "b"],
         ground_with_google_search=True,
-        model_params={"param1": "value1"},
         request_type="TYPE",
     )
     snapshot.assert_match(sql, "generate_text_model_with_options.sql")