From 122842974d582ad40bc3013e9147a47c9d1f7700 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Fri, 9 May 2025 21:22:19 +0000 Subject: [PATCH] test: add gemini test who are ready in production code --- bigframes/ml/llm.py | 1 + tests/system/load/test_llm.py | 3 ++- tests/system/small/ml/test_multimodal_llm.py | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index f93ba8b720..ef74c4ac55 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -83,6 +83,7 @@ _GEMINI_1P5_FLASH_002_ENDPOINT, _GEMINI_2_FLASH_EXP_ENDPOINT, _GEMINI_2_FLASH_001_ENDPOINT, + _GEMINI_2_FLASH_LITE_001_ENDPOINT, ) _CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet" diff --git a/tests/system/load/test_llm.py b/tests/system/load/test_llm.py index abb199b8ab..d346d109be 100644 --- a/tests/system/load/test_llm.py +++ b/tests/system/load/test_llm.py @@ -25,7 +25,7 @@ def llm_remote_text_pandas_df(): return pd.DataFrame( { "prompt": [ - "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", + "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", "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", "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", ], @@ -43,6 +43,7 @@ def llm_remote_text_df(session, llm_remote_text_pandas_df): ( "gemini-1.5-pro-002", "gemini-1.5-flash-002", + "gemini-2.0-flash-001", ), ) def test_llm_gemini_configure_fit( diff --git a/tests/system/small/ml/test_multimodal_llm.py b/tests/system/small/ml/test_multimodal_llm.py index 19ec3d7e14..e066d00cf5 100644 --- a/tests/system/small/ml/test_multimodal_llm.py +++ b/tests/system/small/ml/test_multimodal_llm.py @@ -47,6 +47,7 @@ def test_multimodal_embedding_generator_predict_default_params_success( "gemini-1.5-flash-002", "gemini-2.0-flash-exp", "gemini-2.0-flash-001", + "gemini-2.0-flash-lite-001", ), ) @pytest.mark.flaky(retries=2)