FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

chore: Move long-running LLM tests to the large directory (#2323) · googleapis/google-cloud-python@45462a4 · GitHub

/ google-cloud-python Public

Commit 45462a4

Browse files
google-labs-jules[bot]
authored
chore: Move long-running LLM tests to the large directory (#2323)
Moved long-running ML tests to the large directory to improve test suite organization and execution time management. This involves creating new test files in `tests/system/large/ml/` and appending to existing ones, while removing the corresponding tests from `tests/system/small/ml/`. --- *PR created automatically by Jules for task [11110816135219701367](https://jules.google.com/task/11110816135219701367) started by @tswast* --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
1 parent e58cd20 commit 45462a4

6 files changed

Lines changed: 314 additions & 275 deletions

File tree

‎packages/bigframes/tests/system/large/ml/test_linear_model.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -452,3 +452,39 @@ def test_model_centroids_with_custom_index(penguins_df_default_index):
452452

453453
# If this line executes without errors, the model has correctly ignored the custom index columns
454454
model.predict(X_train.reset_index(drop=True))
455+
456+
457+
def test_linear_reg_model_global_explain(
458+
penguins_linear_model_w_global_explain, new_penguins_df
459+
):
460+
training_data = new_penguins_df.dropna(subset=["body_mass_g"])
461+
X = training_data.drop(columns=["body_mass_g"])
462+
y = training_data[["body_mass_g"]]
463+
penguins_linear_model_w_global_explain.fit(X, y)
464+
global_ex = penguins_linear_model_w_global_explain.global_explain()
465+
assert global_ex.shape == (6, 1)
466+
expected_columns = pd.Index(["attribution"])
467+
pd.testing.assert_index_equal(global_ex.columns, expected_columns)
468+
result = global_ex.to_pandas().drop(["attribution"], axis=1).sort_index()
469+
expected_feature = (
470+
pd.DataFrame(
471+
{
472+
"feature": [
473+
"island",
474+
"species",
475+
"sex",
476+
"flipper_length_mm",
477+
"culmen_depth_mm",
478+
"culmen_length_mm",
479+
]
480+
},
481+
)
482+
.set_index("feature")
483+
.sort_index()
484+
)
485+
pd.testing.assert_frame_equal(
486+
result,
487+
expected_feature,
488+
check_exact=False,
489+
check_index_type=False,
490+
)

‎packages/bigframes/tests/system/large/ml/test_llm.py‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pandas as pd
16+
import pyarrow as pa
17+
import pytest
18+
19+
from bigframes.ml import llm
20+
import bigframes.pandas as bpd
21+
from bigframes.testing import utils
22+
23+
24+
@pytest.mark.parametrize(
25+
"model_name",
26+
(
27+
"gemini-2.0-flash-exp",
28+
"gemini-2.0-flash-001",
29+
"gemini-2.0-flash-lite-001",
30+
"gemini-2.5-pro",
31+
"gemini-2.5-flash",
32+
"gemini-2.5-flash-lite",
33+
),
34+
)
35+
@pytest.mark.flaky(
36+
retries=2
37+
) # usually create model shouldn't be flaky, but this one due to the limited quota of gemini-2.0-flash-exp.
38+
def test_create_load_gemini_text_generator_model(
39+
dataset_id, model_name, session, bq_connection
40+
):
41+
gemini_text_generator_model = llm.GeminiTextGenerator(
42+
model_name=model_name, connection_name=bq_connection, session=session
43+
)
44+
assert gemini_text_generator_model is not None
45+
assert gemini_text_generator_model._bqml_model is not None
46+
47+
# save, load to ensure configuration was kept
48+
reloaded_model = gemini_text_generator_model.to_gbq(
49+
f"{dataset_id}.temp_text_model", replace=True
50+
)
51+
assert f"{dataset_id}.temp_text_model" == reloaded_model._bqml_model.model_name
52+
assert reloaded_model.connection_name == bq_connection
53+
assert reloaded_model.model_name == model_name
54+
55+
56+
@pytest.mark.parametrize(
57+
"model_name",
58+
(
59+
"gemini-2.0-flash-exp",
60+
"gemini-2.0-flash-001",
61+
"gemini-2.0-flash-lite-001",
62+
"gemini-2.5-pro",
63+
"gemini-2.5-flash",
64+
"gemini-2.5-flash-lite",
65+
),
66+
)
67+
# @pytest.mark.flaky(retries=2)
68+
def test_gemini_text_generator_predict_default_params_success(
69+
llm_text_df, model_name, session, bq_connection
70+
):
71+
gemini_text_generator_model = llm.GeminiTextGenerator(
72+
model_name=model_name, connection_name=bq_connection, session=session
73+
)
74+
df = gemini_text_generator_model.predict(llm_text_df).to_pandas()
75+
utils.check_pandas_df_schema_and_index(
76+
df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False
77+
)
78+
79+
80+
@pytest.mark.parametrize(
81+
"model_name",
82+
(
83+
"gemini-2.0-flash-exp",
84+
"gemini-2.0-flash-001",
85+
"gemini-2.0-flash-lite-001",
86+
"gemini-2.5-pro",
87+
"gemini-2.5-flash",
88+
"gemini-2.5-flash-lite",
89+
),
90+
)
91+
@pytest.mark.flaky(retries=2)
92+
def test_gemini_text_generator_predict_with_params_success(
93+
llm_text_df, model_name, session, bq_connection
94+
):
95+
gemini_text_generator_model = llm.GeminiTextGenerator(
96+
model_name=model_name, connection_name=bq_connection, session=session
97+
)
98+
df = gemini_text_generator_model.predict(
99+
llm_text_df, temperature=0.5, max_output_tokens=100, top_k=20, top_p=0.5
100+
).to_pandas()
101+
utils.check_pandas_df_schema_and_index(
102+
df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False
103+
)
104+
105+
106+
@pytest.mark.parametrize(
107+
"model_name",
108+
(
109+
"gemini-2.0-flash-exp",
110+
"gemini-2.0-flash-001",
111+
"gemini-2.0-flash-lite-001",
112+
"gemini-2.5-pro",
113+
"gemini-2.5-flash",
114+
"gemini-2.5-flash-lite",
115+
),
116+
)
117+
@pytest.mark.flaky(retries=2)
118+
def test_gemini_text_generator_multi_cols_predict_success(
119+
llm_text_df: bpd.DataFrame, model_name, session, bq_connection
120+
):
121+
df = llm_text_df.assign(additional_col=1)
122+
gemini_text_generator_model = llm.GeminiTextGenerator(
123+
model_name=model_name, connection_name=bq_connection, session=session
124+
)
125+
pd_df = gemini_text_generator_model.predict(df).to_pandas()
126+
utils.check_pandas_df_schema_and_index(
127+
pd_df,
128+
columns=utils.ML_GENERATE_TEXT_OUTPUT + ["additional_col"],
129+
index=3,
130+
col_exact=False,
131+
)
132+
133+
134+
@pytest.mark.parametrize(
135+
"model_name",
136+
(
137+
"gemini-2.0-flash-exp",
138+
"gemini-2.0-flash-001",
139+
"gemini-2.0-flash-lite-001",
140+
"gemini-2.5-pro",
141+
"gemini-2.5-flash",
142+
"gemini-2.5-flash-lite",
143+
),
144+
)
145+
@pytest.mark.flaky(retries=2)
146+
def test_gemini_text_generator_predict_output_schema_success(
147+
llm_text_df: bpd.DataFrame, model_name, session, bq_connection
148+
):
149+
gemini_text_generator_model = llm.GeminiTextGenerator(
150+
model_name=model_name, connection_name=bq_connection, session=session
151+
)
152+
output_schema = {
153+
"bool_output": "bool",
154+
"int_output": "int64",
155+
"float_output": "float64",
156+
"str_output": "string",
157+
"array_output": "array<int64>",
158+
"struct_output": "struct<number int64>",
159+
}
160+
df = gemini_text_generator_model.predict(llm_text_df, output_schema=output_schema)
161+
assert df["bool_output"].dtype == pd.BooleanDtype()
162+
assert df["int_output"].dtype == pd.Int64Dtype()
163+
assert df["float_output"].dtype == pd.Float64Dtype()
164+
assert df["str_output"].dtype == pd.StringDtype(storage="pyarrow")
165+
assert df["array_output"].dtype == pd.ArrowDtype(pa.list_(pa.int64()))
166+
assert df["struct_output"].dtype == pd.ArrowDtype(
167+
pa.struct([("number", pa.int64())])
168+
)
169+
170+
pd_df = df.to_pandas()
171+
utils.check_pandas_df_schema_and_index(
172+
pd_df,
173+
columns=list(output_schema.keys()) + ["prompt", "full_response", "status"],
174+
index=3,
175+
col_exact=False,
176+
)
177+
178+
179+
@pytest.mark.flaky(retries=2)
180+
@pytest.mark.parametrize(
181+
"model_name",
182+
(
183+
"gemini-2.0-flash-001",
184+
"gemini-2.0-flash-lite-001",
185+
),
186+
)
187+
def test_llm_gemini_score(llm_fine_tune_df_default_index, model_name):
188+
model = llm.GeminiTextGenerator(model_name=model_name)
189+
190+
# Check score to ensure the model was fitted
191+
score_result = model.score(
192+
X=llm_fine_tune_df_default_index[["prompt"]],
193+
y=llm_fine_tune_df_default_index[["label"]],
194+
).to_pandas()
195+
utils.check_pandas_df_schema_and_index(
196+
score_result,
197+
columns=[
198+
"bleu4_score",
199+
"rouge-l_precision",
200+
"rouge-l_recall",
201+
"rouge-l_f1_score",
202+
"evaluation_status",
203+
],
204+
index=1,
205+
)
206+
207+
208+
@pytest.mark.parametrize(
209+
"model_name",
210+
(
211+
"gemini-2.0-flash-001",
212+
"gemini-2.0-flash-lite-001",
213+
),
214+
)
215+
def test_llm_gemini_pro_score_params(llm_fine_tune_df_default_index, model_name):
216+
model = llm.GeminiTextGenerator(model_name=model_name)
217+
218+
# Check score to ensure the model was fitted
219+
score_result = model.score(
220+
X=llm_fine_tune_df_default_index["prompt"],
221+
y=llm_fine_tune_df_default_index["label"],
222+
task_type="classification",
223+
).to_pandas()
224+
utils.check_pandas_df_schema_and_index(
225+
score_result,
226+
columns=[
227+
"precision",
228+
"recall",
229+
"f1_score",
230+
"label",
231+
"evaluation_status",
232+
],
233+
)

‎packages/bigframes/tests/system/large/ml/test_multimodal_llm.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pytest
16+
17+
from bigframes.ml import llm
18+
import bigframes.pandas as bpd
19+
from bigframes.testing import utils
20+
21+
22+
@pytest.mark.parametrize(
23+
"model_name",
24+
(
25+
"gemini-2.0-flash-exp",
26+
"gemini-2.0-flash-001",
27+
"gemini-2.0-flash-lite-001",
28+
),
29+
)
30+
@pytest.mark.flaky(retries=2)
31+
def test_gemini_text_generator_multimodal_input(
32+
images_mm_df: bpd.DataFrame, model_name, session, bq_connection
33+
):
34+
gemini_text_generator_model = llm.GeminiTextGenerator(
35+
model_name=model_name, connection_name=bq_connection, session=session
36+
)
37+
pd_df = gemini_text_generator_model.predict(
38+
images_mm_df, prompt=["Describe", images_mm_df["blob_col"]]
39+
).to_pandas()
40+
utils.check_pandas_df_schema_and_index(
41+
pd_df,
42+
columns=utils.ML_GENERATE_TEXT_OUTPUT + ["blob_col"],
43+
index=2,
44+
col_exact=False,
45+
)

‎packages/bigframes/tests/system/small/ml/test_linear_model.py‎

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -228,42 +228,6 @@ def test_to_gbq_saved_linear_reg_model_scores(
228228
)
229229

230230

231-
def test_linear_reg_model_global_explain(
232-
penguins_linear_model_w_global_explain, new_penguins_df
233-
):
234-
training_data = new_penguins_df.dropna(subset=["body_mass_g"])
235-
X = training_data.drop(columns=["body_mass_g"])
236-
y = training_data[["body_mass_g"]]
237-
penguins_linear_model_w_global_explain.fit(X, y)
238-
global_ex = penguins_linear_model_w_global_explain.global_explain()
239-
assert global_ex.shape == (6, 1)
240-
expected_columns = pandas.Index(["attribution"])
241-
pandas.testing.assert_index_equal(global_ex.columns, expected_columns)
242-
result = global_ex.to_pandas().drop(["attribution"], axis=1).sort_index()
243-
expected_feature = (
244-
pandas.DataFrame(
245-
{
246-
"feature": [
247-
"island",
248-
"species",
249-
"sex",
250-
"flipper_length_mm",
251-
"culmen_depth_mm",
252-
"culmen_length_mm",
253-
]
254-
},
255-
)
256-
.set_index("feature")
257-
.sort_index()
258-
)
259-
pandas.testing.assert_frame_equal(
260-
result,
261-
expected_feature,
262-
check_exact=False,
263-
check_index_type=False,
264-
)
265-
266-
267231
def test_to_gbq_replace(penguins_linear_model, table_id_unique):
268232
penguins_linear_model.to_gbq(table_id_unique, replace=True)
269233
with pytest.raises(google.api_core.exceptions.Conflict):

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL