model_garden/notebooks/official/generative_ai/text_embedding_new_api.ipynb

In [ ]:

# 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
#
#     https://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.

Text Embedding New API

Open in Colab

Open in Colab Enterprise

Open in Workbench

View on GitHub

Overview

This notebook is a code example for how to call our newly released text emebedding models (text-embedding-004, text-multilingual-embedding-002, text-embedding-preview-0815).

Learn more about text embedding api.

Objective

In this tutorial, you learn how to call text embedding latest APIs on two new GA models text-embedding-004, text-multilingual-embedding-002 and one preview model text-embedding-preview-0815.

This tutorial uses the following Google Cloud ML services and resources:

• Vertex LLM SDK

The steps performed include:

• Installation and imports
• Generate embeddings

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI

Learn about Vertex AI pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.

Get Started

Install Vertex AI SDK for python and other required packages

In [ ]:

%pip install --quiet --upgrade google-cloud-aiplatform

Authenticate your notebook environment (Colab only)

Authenticate your environment on Google Colab.

In [ ]:

import sys

if "google.colab" in sys.modules:

    from google.colab import auth

    auth.authenticate_user()

Set Google Cloud project information and initialize Vertex AI SDK for python

To get started using Vertex AI, you must have an existing Google Cloud project and enable the Vertex AI API. Learn more about setting up a project and a development environment.

In [ ]:

PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
LOCATION = "us-central1"  # @param {type:"string"}


from google.cloud import aiplatform

aiplatform.init(project=PROJECT_ID, location=LOCATION)

Import libraries

In [ ]:

from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel

Generate embeddings

1. Set the model name. The latest models are

   • "text-embedding-004" for English.
   • "text-multilingual-embedding-002" for i18n.
   • "text-embedding-preview-0815" for English with text and code embeddings for Python and Java.

   See the language coverage for the list of supported languages.

   See the public doc for the full list of supported models.

2. Set the task type, text, and title (optional and valid only for task type "RETRIEVAL_DOCUMENT"). The valid task types are:

   • "RETRIEVAL_QUERY"
   • "RETRIEVAL_DOCUMENT"
   • "SEMANTIC_SIMILARITY"
   • "CLASSIFICATION"
   • "CLUSTERING"
   • "QUESTION_ANSWERING" (valid only for the latest models)
   • "FACT_VERIFICATION" (valid only for the latest models)
   • "CODE_RETRIEVAL_QUERY" (valid only for the preview-0815 model)

3. Set the output dimensionality (optional and valid only for the latest models).

In [ ]:

# @title { run: "auto" }
MODEL = "text-embedding-004"  # @param ["text-embedding-004", "text-multilingual-embedding-002","text-embedding-preview-0815","text-embedding-preview-0409", "text-multilingual-embedding-preview-0409", "textembedding-gecko@003", "textembedding-gecko-multilingual@001"]
TASK = "RETRIEVAL_DOCUMENT"  # @param ["RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY"]
TEXT = "Banana Muffin?"  # @param {type:"string"}
TITLE = ""  # @param {type:"string"}
OUTPUT_DIMENSIONALITY = 256  # @param [1, 768, "None"] {type:"raw", allow-input:true}

if not MODEL:
    raise ValueError("MODEL must be specified.")
if not TEXT:
    raise ValueError("TEXT must be specified.")
if TITLE and TASK != "RETRIEVAL_DOCUMENT":
    raise ValueError("TITLE can only be specified for TASK 'RETRIEVAL_DOCUMENT'")
if OUTPUT_DIMENSIONALITY is not None and MODEL not in [
    "text-embedding-004",
    "text-multilingual-embedding-002",
    "text-embedding-preview-0815",
    "text-embedding-preview-0409",
    "text-multilingual-embedding-preview-0409",
]:
    raise ValueError(f"OUTPUT_DIMENTIONALITY cannot be specified for model '{MODEL}'.")
if TASK in ["QUESTION_ANSWERING", "FACT_VERIFICATION"] and MODEL not in [
    "text-embedding-004",
    "text-multilingual-embedding-002",
    "text-embedding-preview-0815",
    "text-embedding-preview-0409",
    "text-multilingual-embedding-preview-0409",
]:
    raise ValueError(f"TASK '{TASK}' is not valid for model '{MODEL}'.")
if TASK in ["CODE_RETRIEVAL_QUERY"] and MODEL not in [
    "text-embedding-preview-0815",
]:
    raise ValueError(f"TASK '{TASK}' is not valid for model '{MODEL}'.")

In [ ]:

def embed_text(
    model_name: str,
    task_type: str,
    text: str,
    title: str = "",
    output_dimensionality=None,
) -> list:
    """Generates a text embedding with a Large Language Model."""
    model = TextEmbeddingModel.from_pretrained(model_name)
    text_embedding_input = TextEmbeddingInput(
        task_type=task_type, title=title, text=text
    )
    kwargs = (
        dict(output_dimensionality=output_dimensionality)
        if output_dimensionality
        else {}
    )
    embeddings = model.get_embeddings([text_embedding_input], **kwargs)
    return embeddings[0].values


# Get a text embedding for a downstream task.
embedding = embed_text(
    model_name=MODEL,
    task_type=TASK,
    text=TEXT,
    title=TITLE,
    output_dimensionality=OUTPUT_DIMENSIONALITY,
)
print(len(embedding))  # Expected value: {OUTPUT_DIMENSIONALITY}.

Cleaning up

To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.