From 6c80e61307993dee5fc89c3eb5272bba7a5b39d5 Mon Sep 17 00:00:00 2001
From: Hiroyuki Otomo <otomo_hiroyuki@cyberagent.co.jp>
Date: Fri, 7 Feb 2025 10:40:15 +0900
Subject: [PATCH] chore: add gemini-2.0-flash to the list of supported
 tokenizers.

---
 tests/system/vertexai/test_tokenization.py  | 1 +
 vertexai/tokenization/_tokenizer_loading.py | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

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 = {