Skip to content

Commit 767ca6b

Browse files
shuoweilshuoweil
committed
use fixture for bqml_gemini_text_generator: llm.GeminiTextGenerator,
 bqml_claude3_text_generator: llm.Claude3TextGenerator
1 parent b636ca0 commit 767ca6b

File tree

3 files changed

+70
-49
lines changed

3 files changed

+70
-49
lines changed

‎notebooks/generative_ai/bq_dataframes_llm_gemini_2.ipynb‎

Lines changed: 27 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
},
4545
{
4646
"cell_type": "code",
47-
"execution_count": 3,
47+
"execution_count": 2,
4848
"metadata": {},
4949
"outputs": [],
5050
"source": [
@@ -61,14 +61,14 @@
6161
},
6262
{
6363
"cell_type": "code",
64-
"execution_count": 4,
64+
"execution_count": 3,
6565
"metadata": {},
6666
"outputs": [
6767
{
6868
"name": "stderr",
6969
"output_type": "stream",
7070
"text": [
71-
"/tmp/ipykernel_604828/1304092015.py:1: ApiDeprecationWarning: gemini-pro and gemini-1.5-X are going to be deprecated. Use gemini-2.0-X (https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) instead. \n",
71+
"/tmp/ipykernel_610113/1304092015.py:1: ApiDeprecationWarning: gemini-pro and gemini-1.5-X are going to be deprecated. Use gemini-2.0-X (https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) instead. \n",
7272
" model = llm.GeminiTextGenerator(model_name=\"gemini-1.5-flash-002\")\n",
7373
"/usr/local/google/home/shuowei/src/python-bigquery-dataframes/bigframes/ml/llm.py:981: DefaultLocationWarning: No explicit location is set, so using location US for the session.\n",
7474
" self.session = session or global_session.get_global_session()\n"
@@ -77,7 +77,7 @@
7777
{
7878
"data": {
7979
"text/html": [
80-
"Query job 3f32129b-5f71-4c08-b572-6c7eed201877 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:3f32129b-5f71-4c08-b572-6c7eed201877&page=queryresults\">Open Job</a>"
80+
"Query job f9b52257-1c5a-4bb9-a619-60d7cb5cee79 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:f9b52257-1c5a-4bb9-a619-60d7cb5cee79&page=queryresults\">Open Job</a>"
8181
],
8282
"text/plain": [
8383
"<IPython.core.display.HTML object>"
@@ -89,7 +89,7 @@
8989
{
9090
"data": {
9191
"text/html": [
92-
"Query job 5fb1fa6f-32aa-43c8-b002-c0c21caaaaca is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:5fb1fa6f-32aa-43c8-b002-c0c21caaaaca&page=queryresults\">Open Job</a>"
92+
"Query job 592a26f7-e355-4e67-947b-f5614c34690e is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:592a26f7-e355-4e67-947b-f5614c34690e&page=queryresults\">Open Job</a>"
9393
],
9494
"text/plain": [
9595
"<IPython.core.display.HTML object>"
@@ -112,13 +112,13 @@
112112
},
113113
{
114114
"cell_type": "code",
115-
"execution_count": 5,
115+
"execution_count": 4,
116116
"metadata": {},
117117
"outputs": [
118118
{
119119
"data": {
120120
"text/html": [
121-
"Query job 2276ea5b-2e08-4ed6-af34-49a7d165d145 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:2276ea5b-2e08-4ed6-af34-49a7d165d145&page=queryresults\">Open Job</a>"
121+
"Query job 6cdfea0f-1957-4cd0-a358-d3846b2ddee1 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:6cdfea0f-1957-4cd0-a358-d3846b2ddee1&page=queryresults\">Open Job</a>"
122122
],
123123
"text/plain": [
124124
"<IPython.core.display.HTML object>"
@@ -168,7 +168,7 @@
168168
"[1 rows x 1 columns]"
169169
]
170170
},
171-
"execution_count": 5,
171+
"execution_count": 4,
172172
"metadata": {},
173173
"output_type": "execute_result"
174174
}
@@ -187,13 +187,13 @@
187187
},
188188
{
189189
"cell_type": "code",
190-
"execution_count": 6,
190+
"execution_count": 5,
191191
"metadata": {},
192192
"outputs": [
193193
{
194194
"data": {
195195
"text/html": [
196-
"Query job 9ba21e96-6023-491e-8e83-f2e6fa7df0e7 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:9ba21e96-6023-491e-8e83-f2e6fa7df0e7&page=queryresults\">Open Job</a>"
196+
"Query job 06ce82fd-3041-42ec-848a-92757aece70f is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:06ce82fd-3041-42ec-848a-92757aece70f&page=queryresults\">Open Job</a>"
197197
],
198198
"text/plain": [
199199
"<IPython.core.display.HTML object>"
@@ -206,14 +206,15 @@
206206
"name": "stderr",
207207
"output_type": "stream",
208208
"text": [
209-
"/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/__init__.py:109: PreviewWarning: Interpreting JSON column(s) as StringDtype. This behavior may change in future versions.\n",
210-
" warnings.warn(\n"
209+
"/usr/local/google/home/shuowei/src/python-bigquery-dataframes/bigframes/core/array_value.py:114: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n",
210+
"`db_dtypes` is a preview feature and subject to change.\n",
211+
" warnings.warn(msg, bfe.PreviewWarning)\n"
211212
]
212213
},
213214
{
214215
"data": {
215216
"text/html": [
216-
"Query job 933d45cc-4bc0-4bdf-b4b8-573da2d58be3 is DONE. 2 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:933d45cc-4bc0-4bdf-b4b8-573da2d58be3&page=queryresults\">Open Job</a>"
217+
"Query job 0b7f4fad-4793-4187-866a-ad45d1b8b49c is DONE. 606 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:0b7f4fad-4793-4187-866a-ad45d1b8b49c&page=queryresults\">Open Job</a>"
217218
],
218219
"text/plain": [
219220
"<IPython.core.display.HTML object>"
@@ -225,7 +226,7 @@
225226
{
226227
"data": {
227228
"text/html": [
228-
"Query job 3dda9bc6-84b1-4f4a-8891-85d25d8848ce is DONE. 4.3 kB processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:3dda9bc6-84b1-4f4a-8891-85d25d8848ce&page=queryresults\">Open Job</a>"
229+
"Query job 5774da7a-23cb-45e8-bf6a-87d7ee2ebacd is DONE. 614 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:5774da7a-23cb-45e8-bf6a-87d7ee2ebacd&page=queryresults\">Open Job</a>"
229230
],
230231
"text/plain": [
231232
"<IPython.core.display.HTML object>"
@@ -264,7 +265,7 @@
264265
" <tbody>\n",
265266
" <tr>\n",
266267
" <th>0</th>\n",
267-
" <td>Alright, let's talk about Gemini 2.0! It's a b...</td>\n",
268+
" <td>There's no officially recognized technology or...</td>\n",
268269
" <td>&lt;NA&gt;</td>\n",
269270
" <td></td>\n",
270271
" <td>Tell me something about Gemini 2.0.</td>\n",
@@ -276,7 +277,7 @@
276277
],
277278
"text/plain": [
278279
" ml_generate_text_llm_result \\\n",
279-
"0 Alright, let's talk about Gemini 2.0! It's a b... \n",
280+
"0 There's no officially recognized technology or... \n",
280281
"\n",
281282
" ml_generate_text_rai_result ml_generate_text_status \\\n",
282283
"0 <NA> \n",
@@ -287,7 +288,7 @@
287288
"[1 rows x 4 columns]"
288289
]
289290
},
290-
"execution_count": 6,
291+
"execution_count": 5,
291292
"metadata": {},
292293
"output_type": "execute_result"
293294
}
@@ -306,13 +307,13 @@
306307
},
307308
{
308309
"cell_type": "code",
309-
"execution_count": 7,
310+
"execution_count": 6,
310311
"metadata": {},
311312
"outputs": [
312313
{
313314
"data": {
314315
"text/html": [
315-
"Copy job 8e68af62-e7ab-475b-99c9-b79e8ba3c40b is DONE. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:8e68af62-e7ab-475b-99c9-b79e8ba3c40b&page=queryresults\">Open Job</a>"
316+
"Copy job 8fe43863-54be-4c8a-b2b3-5fdba16a1bc0 is DONE. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:8fe43863-54be-4c8a-b2b3-5fdba16a1bc0&page=queryresults\">Open Job</a>"
316317
],
317318
"text/plain": [
318319
"<IPython.core.display.HTML object>"
@@ -325,17 +326,14 @@
325326
"name": "stderr",
326327
"output_type": "stream",
327328
"text": [
328-
"/usr/local/google/home/garrettwu/src/bigframes/bigframes/ml/llm.py:803: PreviewWarning: Model gemini-2.0-flash-exp is subject to the \"Pre-GA Offerings Terms\" in the General Service Terms section of the\n",
329-
" Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available \"as is\"\n",
330-
" and might have limited support. For more information, see the launch stage descriptions\n",
331-
" (https://cloud.google.com/products#product-launch-stages).\n",
332-
" warnings.warn(\n"
329+
"/usr/local/google/home/shuowei/src/python-bigquery-dataframes/bigframes/ml/llm.py:1023: ApiDeprecationWarning: gemini-pro and gemini-1.5-X are going to be deprecated. Use gemini-2.0-X (https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) instead. \n",
330+
" model = cls(\n"
333331
]
334332
},
335333
{
336334
"data": {
337335
"text/html": [
338-
"Query job cae7f929-d8cb-4819-a644-ac832cdc0912 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:cae7f929-d8cb-4819-a644-ac832cdc0912&page=queryresults\">Open Job</a>"
336+
"Query job eed27608-273f-45b2-bb2e-12e1b42c5c04 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:eed27608-273f-45b2-bb2e-12e1b42c5c04&page=queryresults\">Open Job</a>"
339337
],
340338
"text/plain": [
341339
"<IPython.core.display.HTML object>"
@@ -347,12 +345,12 @@
347345
{
348346
"data": {
349347
"text/plain": [
350-
"GeminiTextGenerator(connection_name='bigframes-dev.us.bigframes-rf-connection',\n",
351-
" model_name='gemini-2.0-flash-exp',\n",
352-
" session=<bigframes.session.Session object at 0x7f8210630370>)"
348+
"GeminiTextGenerator(connection_name='bigframes-dev.US.bigframes-default-connection',\n",
349+
" model_name='gemini-1.5-flash-002',\n",
350+
" session=<bigframes.session.Session object at 0x7f38e5eab7c0>)"
353351
]
354352
},
355-
"execution_count": 7,
353+
"execution_count": 6,
356354
"metadata": {},
357355
"output_type": "execute_result"
358356
}

‎tests/system/small/ml/conftest.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -397,3 +397,20 @@ def imported_xgboost_model(
397397
output={"predicted_label": "float64"},
398398
model_path=imported_xgboost_array_model_path,
399399
)
400+
401+
402+
@pytest.fixture(scope="session")
403+
def bqml_gemini_text_generator(bq_connection, session) -> llm.GeminiTextGenerator:
404+
return llm.GeminiTextGenerator(
405+
model_name="gemini-1.5-flash-002", # Or your desired Gemini model
406+
connection_name=bq_connection,
407+
session=session,
408+
)
409+
410+
411+
@pytest.fixture(scope="session")
412+
def bqml_claude3_text_generator(bq_connection, session) -> llm.Claude3TextGenerator:
413+
return llm.Claude3TextGenerator(
414+
connection_name=bq_connection,
415+
session=session,
416+
)

‎tests/system/small/ml/test_llm.py‎

Lines changed: 26 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -458,16 +458,12 @@ def __eq__(self, other):
458458

459459
@pytest.mark.parametrize(
460460
(
461-
"text_generator_model",
461+
"model_class",
462462
"options",
463463
),
464464
[
465465
(
466-
llm.GeminiTextGenerator(
467-
connection_name="test_connection",
468-
session=mock.Mock(),
469-
model_name="gemini-1.5-flash-002",
470-
),
466+
llm.GeminiTextGenerator,
471467
{
472468
"temperature": 0.9,
473469
"max_output_tokens": 8192,
@@ -478,9 +474,7 @@ def __eq__(self, other):
478474
},
479475
),
480476
(
481-
llm.Claude3TextGenerator(
482-
connection_name="test_connection", session=mock.Mock()
483-
),
477+
llm.Claude3TextGenerator,
484478
{
485479
"max_output_tokens": 128,
486480
"top_k": 40,
@@ -491,7 +485,11 @@ def __eq__(self, other):
491485
],
492486
)
493487
def test_text_generator_retry_success(
494-
session, bq_connection, text_generator_model, options
488+
session,
489+
model_name,
490+
options,
491+
bqml_gemini_text_generator: llm.GeminiTextGenerator,
492+
bqml_claude3_text_generator: llm.Claude3TextGenerator,
495493
):
496494
# Requests.
497495
df0 = EqCmpAllDataFrame(
@@ -567,6 +565,11 @@ def test_text_generator_retry_success(
567565
),
568566
]
569567

568+
text_generator_model = (
569+
bqml_gemini_text_generator
570+
if (model_name == llm.GeminiTextGenerator)
571+
else bqml_claude3_text_generator
572+
)
570573
text_generator_model._bqml_model = mock_bqml_model
571574

572575
# 3rd retry isn't triggered
@@ -599,16 +602,12 @@ def test_text_generator_retry_success(
599602

600603
@pytest.mark.parametrize(
601604
(
602-
"text_generator_model",
605+
"model_name",
603606
"options",
604607
),
605608
[
606609
(
607-
llm.GeminiTextGenerator(
608-
connection_name="test_connection",
609-
session=mock.Mock(),
610-
model_name="gemini-1.5-flash-002",
611-
),
610+
llm.GeminiTextGenerator,
612611
{
613612
"temperature": 0.9,
614613
"max_output_tokens": 8192,
@@ -619,9 +618,7 @@ def test_text_generator_retry_success(
619618
},
620619
),
621620
(
622-
llm.Claude3TextGenerator(
623-
connection_name="test_connection", session=mock.Mock()
624-
),
621+
llm.Claude3TextGenerator,
625622
{
626623
"max_output_tokens": 128,
627624
"top_k": 40,
@@ -632,7 +629,11 @@ def test_text_generator_retry_success(
632629
],
633630
)
634631
def test_text_generator_retry_no_progress(
635-
session, bq_connection, text_generator_model, options
632+
session,
633+
model_name,
634+
options,
635+
bqml_gemini_text_generator: llm.GeminiTextGenerator,
636+
bqml_claude3_text_generator: llm.Claude3TextGenerator,
636637
):
637638
# Requests.
638639
df0 = EqCmpAllDataFrame(
@@ -687,6 +688,11 @@ def test_text_generator_retry_no_progress(
687688
),
688689
]
689690

691+
text_generator_model = (
692+
bqml_gemini_text_generator
693+
if (model_name == llm.GeminiTextGenerator)
694+
else bqml_claude3_text_generator
695+
)
690696
text_generator_model._bqml_model = mock_bqml_model
691697

692698
# No progress, only conduct retry once

0 commit comments

Comments
 (0)