Skip to content

Commit de9efdb

Browse files
authored
test: add gemini-2.0-flash-lite-011 test (#1729)
1 parent 55c07e9 commit de9efdb

File tree

2 files changed

+4
-2
lines changed

2 files changed

+4
-2
lines changed

bigframes/ml/llm.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -529,7 +529,8 @@ def fit(
529529
y: utils.ArrayType,
530530
) -> GeminiTextGenerator:
531531
"""Fine tune GeminiTextGenerator model. Only support "gemini-1.5-pro-002",
532-
and "gemini-1.5-flash-002" models for now.
532+
"gemini-1.5-flash-002", "gemini-2.0-flash-001",
533+
and "gemini-2.0-flash-lite-001"models for now.
533534
534535
.. note::
535536
@@ -549,7 +550,7 @@ def fit(
549550
"""
550551
if self.model_name not in _GEMINI_FINE_TUNE_SCORE_ENDPOINTS:
551552
msg = exceptions.format_message(
552-
"fit() only supports gemini-1.5-pro-002, or gemini-1.5-flash-002 model."
553+
"fit() only supports gemini-1.5-pro-002, gemini-1.5-flash-002, gemini-2.0-flash-001, or gemini-2.0-flash-lite-001 model."
553554
)
554555
warnings.warn(msg)
555556

tests/system/load/test_llm.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ def llm_remote_text_df(session, llm_remote_text_pandas_df):
4444
"gemini-1.5-pro-002",
4545
"gemini-1.5-flash-002",
4646
"gemini-2.0-flash-001",
47+
"gemini-2.0-flash-lite-001",
4748
),
4849
)
4950
def test_llm_gemini_configure_fit(

0 commit comments

Comments
 (0)