Adding an example to show how to get the token count (GoogleCloudPlatform#8944)

glaforge · web-flow · commit 7d4f71108145 · 2023-12-20T20:17:52.000+05:30
* Revert "Use best practice in multimodal prompts to put the image first, the string prompt second" This reverts commit 594a259. * Add token count example in Java * Put back the switch in image/prompt order * Put back the switch in image/prompt order * Add missing license header * Fix indentation * Let's please checkstyle * Reduce scioe if region tags * Correct indentation to please checkstyle
diff --git a/vertexai/snippets/src/main/java/vertexai/gemini/GetTokenCount.java b/vertexai/snippets/src/main/java/vertexai/gemini/GetTokenCount.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package vertexai.gemini;
+
+// [START aiplatform_gemini_token_count]
+import com.google.cloud.vertexai.VertexAI;
+import com.google.cloud.vertexai.api.CountTokensResponse;
+import com.google.cloud.vertexai.generativeai.preview.GenerativeModel;
+import java.io.IOException;
+// [END aiplatform_gemini_token_count]
+
+public class GetTokenCount {
+  public static void main(String[] args) throws IOException {
+    // TODO(developer): Replace these variables before running the sample.
+    String projectId = "your-google-cloud-project-id";
+    String location = "us-central1";
+    String modelName = "gemini-pro-vision";
+
+    String textPrompt = "How many tokens are there in this prompt?";
+    getTokenCount(projectId, location, modelName, textPrompt);
+  }
+
+  // [START aiplatform_gemini_token_count]
+  public static int getTokenCount(String projectId, String location, String modelName,
+                                  String textPrompt)
+      throws IOException {
+    try (VertexAI vertexAI = new VertexAI(projectId, location)) {
+      GenerativeModel model = new GenerativeModel(modelName, vertexAI);
+      CountTokensResponse response = model.countTokens(textPrompt);
+
+      int tokenCount = response.getTotalTokens();
+      System.out.println("There are " + tokenCount + " tokens in the prompt.");
+
+      return tokenCount;
+    }
+  }
+  // [END aiplatform_gemini_token_count]
+}
diff --git a/vertexai/snippets/src/main/java/vertexai/gemini/Quickstart.java b/vertexai/snippets/src/main/java/vertexai/gemini/Quickstart.java
@@ -46,8 +46,8 @@ public static String quickstart(String projectId, String location, String modelN
 
       GenerativeModel model = new GenerativeModel(modelName, vertexAI);
       GenerateContentResponse response = model.generateContent(ContentMaker.fromMultiModalData(
-          "What's in this photo",
-          PartMaker.fromMimeTypeAndData("image/jpg", imageUri)
+          PartMaker.fromMimeTypeAndData("image/jpg", imageUri),
+          "What's in this photo"
       ));
 
       return response.toString();
diff --git a/vertexai/snippets/src/test/java/vertexai/gemini/SnippetsIT.java b/vertexai/snippets/src/test/java/vertexai/gemini/SnippetsIT.java
@@ -183,4 +183,13 @@ public void testSafetySettings() throws Exception {
     assertThat(output).isNotEmpty();
     assertThat(output).contains("reasons?");
   }
+
+  @Test
+  public void testTokenCount() throws Exception {
+    String textPrompt = "How many tokens are there in this prompt?";
+
+    int tokenCount = GetTokenCount.getTokenCount(PROJECT_ID, LOCATION, GEMINI_PRO_VISION,
+        textPrompt);
+    assertThat(tokenCount).isGreaterThan(6);
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -183,4 +183,13 @@ public void testSafetySettings() throws Exception {`
`183`	`183`	`assertThat(output).isNotEmpty();`
`184`	`184`	`assertThat(output).contains("reasons?");`
`185`	`185`	`}`
	`186`	`+`
	`187`	`+ @Test`
	`188`	`+ public void testTokenCount() throws Exception {`
	`189`	`+ String textPrompt = "How many tokens are there in this prompt?";`
	`190`	`+`
	`191`	`+ int tokenCount = GetTokenCount.getTokenCount(PROJECT_ID, LOCATION, GEMINI_PRO_VISION,`
	`192`	`+ textPrompt);`
	`193`	`+ assertThat(tokenCount).isGreaterThan(6);`
	`194`	`+ }`
`186`	`195`	`}`