Count tokens for Gemini

The code sample demonstrates how to use the Vertex AI Generative Models API to count the number of tokens in a prompt and generate content using the Gemini model.

Explore further

For detailed documentation that includes this code sample, see the following:

• CountTokens API

Code sample

Go

Before trying this sample, follow the Go setup instructions in the Vertex AI quickstart using client libraries. For more information, see the Vertex AI Go API reference documentation.

To authenticate to Vertex AI, set up Application Default Credentials. For more information, see Set up authentication for a local development environment.

import (
	"context"
	"fmt"
	"io"
	"mime"
	"path/filepath"

	"cloud.google.com/go/vertexai/genai"
)

// countTokensMultimodal finds the number of tokens for a multimodal prompt (video+text), and writes to w. Then,
// it calls the model with the multimodal prompt and writes token counts from the response metadata to w.
//
// video is a Google Cloud Storage path starting with "gs://"
func countTokensMultimodal(w io.Writer, projectID, location, modelName string) error {
	// location := "us-central1"
	// modelName := "gemini-1.5-flash-001"
	prompt := "Provide a description of the video."
	video := "gs://cloud-samples-data/generative-ai/video/pixel8.mp4"

	ctx := context.Background()

	client, err := genai.NewClient(ctx, projectID, location)
	if err != nil {
		return fmt.Errorf("unable to create client: %w", err)
	}
	defer client.Close()

	model := client.GenerativeModel(modelName)

	part1 := genai.Text(prompt)

	// Given a video file URL, prepare video file as genai.Part
	part2 := genai.FileData{
		MIMEType: mime.TypeByExtension(filepath.Ext(video)),
		FileURI:  video,
	}

	// Finds the total number of tokens for the 2 parts (text, video) of the multimodal prompt,
	// before actually calling the model for inference.
	resp, err := model.CountTokens(ctx, part1, part2)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Number of tokens for the multimodal video prompt: %d\n", resp.TotalTokens)

	res, err := model.GenerateContent(ctx, part1, part2)
	if err != nil {
		return fmt.Errorf("unable to generate contents: %w", err)
	}

	// The token counts are also provided in the model response metadata, after inference.
	fmt.Fprintln(w, "\nModel response")
	md := res.UsageMetadata
	fmt.Fprintf(w, "Prompt Token Count: %d\n", md.PromptTokenCount)
	fmt.Fprintf(w, "Candidates Token Count: %d\n", md.CandidatesTokenCount)
	fmt.Fprintf(w, "Total Token Count: %d\n", md.TotalTokenCount)

	return nil
}

Python

Before trying this sample, follow the Python setup instructions in the Vertex AI quickstart using client libraries. For more information, see the Vertex AI Python API reference documentation.

To authenticate to Vertex AI, set up Application Default Credentials. For more information, see Set up authentication for a local development environment.

import vertexai
from vertexai.generative_models import GenerativeModel, Part

# TODO(developer): Update and un-comment below lines
# project_id = "PROJECT_ID"

vertexai.init(project=project_id, location="us-central1")

model = GenerativeModel(model_name="gemini-1.5-flash-001")

contents = [
    Part.from_uri(
        "gs://cloud-samples-data/generative-ai/video/pixel8.mp4",
        mime_type="video/mp4",
    ),
    "Provide a description of the video.",
]

# Prompt tokens count
response = model.count_tokens(contents)
print(f"Prompt Token Count: {response.total_tokens}")
print(f"Prompt Character Count: {response.total_billable_characters}")

# Send text to Gemini
response = model.generate_content(contents)
usage_metadata = response.usage_metadata

# Response tokens count
print(f"Prompt Token Count: {usage_metadata.prompt_token_count}")
print(f"Candidates Token Count: {usage_metadata.candidates_token_count}")
print(f"Total Token Count: {usage_metadata.total_token_count}")

What's next

To search and filter code samples for other Google Cloud products, see the Google Cloud sample browser.