diff --git a/tests/system/vertexai/test_tokenization.py b/tests/system/vertexai/test_tokenization.py index 105bcf0a2d..f1387d320c 100644 --- a/tests/system/vertexai/test_tokenization.py +++ b/tests/system/vertexai/test_tokenization.py @@ -44,6 +44,7 @@ "gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.5-flash", + "gemini-2.0-flash", "gemini-1.5-flash-002", "gemini-1.5-pro-002", ] diff --git a/vertexai/tokenization/_tokenizer_loading.py b/vertexai/tokenization/_tokenizer_loading.py index b8167b2b47..342c9db159 100644 --- a/vertexai/tokenization/_tokenizer_loading.py +++ b/vertexai/tokenization/_tokenizer_loading.py @@ -34,7 +34,12 @@ class _TokenizerConfig: _GEMMA_TOKENIZER = "google/gemma" # SoT: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models -_GEMINI_MODEL_NAMES = ["gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.5-flash"] +_GEMINI_MODEL_NAMES = [ + "gemini-1.0-pro", + "gemini-1.5-pro", + "gemini-1.5-flash", + "gemini-2.0-flash", +] _GEMINI_STABLE_MODEL_NAMES = [ "gemini-1.0-pro-001", "gemini-1.0-pro-002", @@ -42,6 +47,7 @@ class _TokenizerConfig: "gemini-1.5-flash-001", "gemini-1.5-flash-002", "gemini-1.5-pro-002", + "gemini-2.0-flash-001", ] _TOKENIZERS = {