using vertex ai api

app.py

@@ -1,26 +1,31 @@
-from vertexai.preview import tokenization
-import gradio as gr
+from google import genai
+import os
 
 color_map = {
-    '0': '#ccbfee',
-    '1': '#beedc6',
-    '2': '#f4d7ab',
-    '3': '#f4aeb1',
-    '4': '#a4dcf3',
+    '3': '#ccbfee',
+    '2': '#beedc6',
+    '1': '#f4d7ab',
+    '0': '#f4aeb1',
+    # '4': '#a4dcf3',
 }
 
-def tokens_len(prompt, model):
-    model_name = "gemini-1.5-pro"
-    tokenizer = tokenization.get_tokenizer_for_model(model)
+client_vertex = genai.Client(
+    vertexai=True,
+    api_key=os.getenv("VERTEX_API_KEY")
+)
 
-    result = tokenizer.compute_tokens(prompt)
+def tokens_len(prompt, model):
+    result = client_vertex.models.compute_tokens(
+        model=model,
+        contents=prompt
+    )
     tokens = result.tokens_info[0].tokens
 
     return (
         f'{len(tokens)} tokens, {len(prompt)} chracters',
         [
             (
-                token.decode('utf-8').replace("\n", "\\n"),
+                token.decode('utf-8').replace("\n", "\\n"), 
                 str(i % len(color_map))
             ) for i, token in enumerate(tokens)
         ]
@@ -35,10 +40,10 @@ web_tokenizer = gr.Interface(
         gr.Dropdown(
             label='模型名稱',
             choices = [
-                'gemini-1.5-pro',
-                'gemini-1.5-flash',
+                'gemini-3-pro-preview',
+                'gemini-2.5-flash',
             ],
-            value='gemini-1.5-pro')
+            value='gemini-3-pro-preview')
     ],
     outputs = [
         gr.Textbox(
@@ -50,4 +55,4 @@ web_tokenizer = gr.Interface(
     ]
 )
 
-web_tokenizer.launch(share=True, debug=True)
+web_tokenizer.launch(share=True, debug=True)