Skip to content

Commit ab24b6a

Browse files
authored
test: add gemini test who are ready in production code (#1718)
1 parent 6b1fe32 commit ab24b6a

File tree

3 files changed

+4
-1
lines changed

3 files changed

+4
-1
lines changed

bigframes/ml/llm.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@
8383
_GEMINI_1P5_FLASH_002_ENDPOINT,
8484
_GEMINI_2_FLASH_EXP_ENDPOINT,
8585
_GEMINI_2_FLASH_001_ENDPOINT,
86+
_GEMINI_2_FLASH_LITE_001_ENDPOINT,
8687
)
8788

8889
_CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"

tests/system/load/test_llm.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ def llm_remote_text_pandas_df():
2525
return pd.DataFrame(
2626
{
2727
"prompt": [
28-
"Please do sentiment analysis on the following text and only output a number from 0 to 5where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i feel beautifully emotional knowing that these women of whom i knew just a handful were holding me and my baba on our journey",
28+
"Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i feel beautifully emotional knowing that these women of whom i knew just a handful were holding me and my baba on our journey",
2929
"Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: i was feeling a little vain when i did this one",
3030
"Please do sentiment analysis on the following text and only output a number from 0 to 5 where 0 means sadness, 1 means joy, 2 means love, 3 means anger, 4 means fear, and 5 means surprise. Text: a father of children killed in an accident",
3131
],
@@ -43,6 +43,7 @@ def llm_remote_text_df(session, llm_remote_text_pandas_df):
4343
(
4444
"gemini-1.5-pro-002",
4545
"gemini-1.5-flash-002",
46+
"gemini-2.0-flash-001",
4647
),
4748
)
4849
def test_llm_gemini_configure_fit(

tests/system/small/ml/test_multimodal_llm.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ def test_multimodal_embedding_generator_predict_default_params_success(
4747
"gemini-1.5-flash-002",
4848
"gemini-2.0-flash-exp",
4949
"gemini-2.0-flash-001",
50+
"gemini-2.0-flash-lite-001",
5051
),
5152
)
5253
@pytest.mark.flaky(retries=2)

0 commit comments

Comments
 (0)