Skip to content

Commit 08c0c0c

Browse files
GarrettWuGarrettWu
authored
docs: update docs and tests for Gemini 2.5 models (#2279)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent bc33c98 commit 08c0c0c

File tree

3 files changed

+37
-5
lines changed

3 files changed

+37
-5
lines changed

‎bigframes/ml/llm.py‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@
5454
_GEMINI_2_FLASH_001_ENDPOINT = "gemini-2.0-flash-001"
5555
_GEMINI_2_FLASH_LITE_001_ENDPOINT = "gemini-2.0-flash-lite-001"
5656
_GEMINI_2P5_PRO_PREVIEW_ENDPOINT = "gemini-2.5-pro-preview-05-06"
57+
_GEMINI_2P5_PRO_ENDPOINT = "gemini-2.5-pro"
58+
_GEMINI_2P5_FLASH_ENDPOINT = "gemini-2.5-flash"
59+
_GEMINI_2P5_FLASH_LITE_ENDPOINT = "gemini-2.5-flash-lite"
60+
5761
_GEMINI_ENDPOINTS = (
5862
_GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
5963
_GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT,
@@ -64,6 +68,9 @@
6468
_GEMINI_2_FLASH_EXP_ENDPOINT,
6569
_GEMINI_2_FLASH_001_ENDPOINT,
6670
_GEMINI_2_FLASH_LITE_001_ENDPOINT,
71+
_GEMINI_2P5_PRO_ENDPOINT,
72+
_GEMINI_2P5_FLASH_ENDPOINT,
73+
_GEMINI_2P5_FLASH_LITE_ENDPOINT,
6774
)
6875
_GEMINI_PREVIEW_ENDPOINTS = (
6976
_GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -84,6 +91,9 @@
8491
_GEMINI_2_FLASH_EXP_ENDPOINT,
8592
_GEMINI_2_FLASH_001_ENDPOINT,
8693
_GEMINI_2_FLASH_LITE_001_ENDPOINT,
94+
_GEMINI_2P5_PRO_ENDPOINT,
95+
_GEMINI_2P5_FLASH_ENDPOINT,
96+
_GEMINI_2P5_FLASH_LITE_ENDPOINT,
8797
)
8898

8999
_CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"
@@ -419,20 +429,21 @@ class GeminiTextGenerator(base.RetriableRemotePredictor):
419429
"""Gemini text generator LLM model.
420430
421431
.. note::
422-
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.
432+
gemini-1.5-X are going to be deprecated. Use gemini-2.5-X (https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) instead.
423433
424434
Args:
425435
model_name (str, Default to "gemini-2.0-flash-001"):
426436
The model for natural language tasks. Accepted values are
427437
"gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514",
428438
"gemini-1.5-pro-001", "gemini-1.5-pro-002", "gemini-1.5-flash-001",
429439
"gemini-1.5-flash-002", "gemini-2.0-flash-exp",
430-
"gemini-2.0-flash-lite-001", and "gemini-2.0-flash-001".
440+
"gemini-2.0-flash-lite-001", "gemini-2.0-flash-001",
441+
"gemini-2.5-pro", "gemini-2.5-flash" and "gemini-2.5-flash-lite".
431442
If no setting is provided, "gemini-2.0-flash-001" will be used by
432443
default and a warning will be issued.
433444
434445
.. note::
435-
"gemini-1.5-X" is going to be deprecated. Please use gemini-2.0-X instead. For example, "gemini-2.0-flash-001".
446+
"gemini-1.5-X" is going to be deprecated. Please use gemini-2.5-X instead. For example, "gemini-2.5-flash".
436447
"gemini-2.0-flash-exp", "gemini-1.5-pro-preview-0514" and "gemini-1.5-flash-preview-0514" is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the
437448
Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
438449
and might have limited support. For more information, see the launch stage descriptions
@@ -462,6 +473,9 @@ def __init__(
462473
"gemini-2.0-flash-exp",
463474
"gemini-2.0-flash-001",
464475
"gemini-2.0-flash-lite-001",
476+
"gemini-2.5-pro",
477+
"gemini-2.5-flash",
478+
"gemini-2.5-flash-lite",
465479
]
466480
] = None,
467481
session: Optional[bigframes.Session] = None,
@@ -510,7 +524,7 @@ def _create_bqml_model(self):
510524
msg = exceptions.format_message(
511525
_MODEL_DEPRECATE_WARNING.format(
512526
model_name=self.model_name,
513-
new_model_name="gemini-2.0-X",
527+
new_model_name="gemini-2.5-X",
514528
link="https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator",
515529
)
516530
)

‎bigframes/ml/loader.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@
6767
llm._GEMINI_2_FLASH_001_ENDPOINT: llm.GeminiTextGenerator,
6868
llm._GEMINI_2_FLASH_LITE_001_ENDPOINT: llm.GeminiTextGenerator,
6969
llm._GEMINI_2P5_PRO_PREVIEW_ENDPOINT: llm.GeminiTextGenerator,
70+
llm._GEMINI_2P5_FLASH_ENDPOINT: llm.GeminiTextGenerator,
71+
llm._GEMINI_2P5_FLASH_LITE_ENDPOINT: llm.GeminiTextGenerator,
72+
llm._GEMINI_2P5_PRO_ENDPOINT: llm.GeminiTextGenerator,
7073
llm._CLAUDE_3_HAIKU_ENDPOINT: llm.Claude3TextGenerator,
7174
llm._CLAUDE_3_SONNET_ENDPOINT: llm.Claude3TextGenerator,
7275
llm._CLAUDE_3_5_SONNET_ENDPOINT: llm.Claude3TextGenerator,

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,9 @@ def test_create_load_multimodal_embedding_generator_model(
111111
"gemini-2.0-flash-exp",
112112
"gemini-2.0-flash-001",
113113
"gemini-2.0-flash-lite-001",
114+
"gemini-2.5-pro",
115+
"gemini-2.5-flash",
116+
"gemini-2.5-flash-lite",
114117
),
115118
)
116119
@pytest.mark.flaky(
@@ -140,9 +143,12 @@ def test_create_load_gemini_text_generator_model(
140143
"gemini-2.0-flash-exp",
141144
"gemini-2.0-flash-001",
142145
"gemini-2.0-flash-lite-001",
146+
"gemini-2.5-pro",
147+
"gemini-2.5-flash",
148+
"gemini-2.5-flash-lite",
143149
),
144150
)
145-
@pytest.mark.flaky(retries=2)
151+
# @pytest.mark.flaky(retries=2)
146152
def test_gemini_text_generator_predict_default_params_success(
147153
llm_text_df, model_name, session, bq_connection
148154
):
@@ -161,6 +167,9 @@ def test_gemini_text_generator_predict_default_params_success(
161167
"gemini-2.0-flash-exp",
162168
"gemini-2.0-flash-001",
163169
"gemini-2.0-flash-lite-001",
170+
"gemini-2.5-pro",
171+
"gemini-2.5-flash",
172+
"gemini-2.5-flash-lite",
164173
),
165174
)
166175
@pytest.mark.flaky(retries=2)
@@ -184,6 +193,9 @@ def test_gemini_text_generator_predict_with_params_success(
184193
"gemini-2.0-flash-exp",
185194
"gemini-2.0-flash-001",
186195
"gemini-2.0-flash-lite-001",
196+
"gemini-2.5-pro",
197+
"gemini-2.5-flash",
198+
"gemini-2.5-flash-lite",
187199
),
188200
)
189201
@pytest.mark.flaky(retries=2)
@@ -209,6 +221,9 @@ def test_gemini_text_generator_multi_cols_predict_success(
209221
"gemini-2.0-flash-exp",
210222
"gemini-2.0-flash-001",
211223
"gemini-2.0-flash-lite-001",
224+
"gemini-2.5-pro",
225+
"gemini-2.5-flash",
226+
"gemini-2.5-flash-lite",
212227
),
213228
)
214229
@pytest.mark.flaky(retries=2)

0 commit comments

Comments
 (0)