model_garden/notebooks/community/cohere/cohere_embedding_with_matching_engine.ipynb

In [ ]:

# Copyright 2021 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.

Run in Google Cloud Notebooks

View on GitHub

Overview

In this notebook, you will learn how to use co.embed to quickly capture semantic information about some input data. You'll then use those new features with Vertex AI Matching Engine's Approximate Nearest Neighbor (ANN) service to find similar texts.

Dataset

The dataset used for this tutorial is the ag_news_subset from TensorFlow Datasets. The final solution will be able to find articles that are similar to a user-provided topic.

Objective

In this notebook, you will learn how to create embeddings with Cohere's API and then leverage Vertex AI Matching Engine to create an Index and query against indexes to find similar texts.

Costs

This tutorial uses billable components of Google Cloud and Cohere:

• Vertex AI
• Cloud Storage
• Cohere co.embed endpoint usage

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

Learn about Cohere Pricing

Before you begin

Lets set the variables that will be used throughout this demo

In [ ]:

COHERE_API_KEY = "{API KEY}"
GOOGLE_PROJECT_ID = "{Project ID}"
NETWORK_NAME = "{Network Name}"
PEERING_RANGE_NAME = "{Range Name}"
BUCKET_NAME = "gs://{Bucket Name}"
REGION = "us-central1"

Creating a VPC Network

• Prepare a VPC network. To reduce any network overhead that might lead to unnecessary increase in overhead latency, it is best to call the ANN endpoints from your VPC via a direct VPC Peering connection. The following section describes how to setup a VPC Peering connection if you don't have one. This is a one-time initial setup task. You can also reuse existing VPC network and skip this section.
• WARNING: The match service gRPC API (to create online queries against your deployed index) has to be executed in a Google Cloud Notebook instance that is created with the following requirements:
  • In the same region as where your ANN service is deployed (for example, if you set REGION = "us-central1" as same as the tutorial, the notebook instance has to be in us-central1).
  • Make sure you select the VPC network you created for ANN service (instead of using the "default" one). That is, you will have to create the VPC network below and then create a new notebook instance that uses that VPC.
  • If you run it in the colab or a Google Cloud Notebook instance in a different VPC network or region, the gRPC API will fail to peer the network (InactiveRPCError).

In [ ]:

PROJECT_ID = GOOGLE_PROJECT_ID  # @param {type:"string"}

# Create a VPC network
! gcloud compute networks create {NETWORK_NAME} --bgp-routing-mode=regional --subnet-mode=auto --project={PROJECT_ID}

# Add necessary firewall rules
! gcloud compute firewall-rules create {NETWORK_NAME}-allow-icmp --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow icmp

! gcloud compute firewall-rules create {NETWORK_NAME}-allow-internal --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow all --source-ranges 10.128.0.0/9

! gcloud compute firewall-rules create {NETWORK_NAME}-allow-rdp --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow tcp:3389

! gcloud compute firewall-rules create {NETWORK_NAME}-allow-ssh --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow tcp:22

# Reserve IP range
! gcloud compute addresses create {PEERING_RANGE_NAME} --global --prefix-length=16 --network={NETWORK_NAME} --purpose=VPC_PEERING --project={PROJECT_ID} --description="peering range for uCAIP Haystack."

# Set up peering with service networking
! gcloud services vpc-peerings connect --service=servicenetworking.googleapis.com --network={NETWORK_NAME} --ranges={PEERING_RANGE_NAME} --project={PROJECT_ID}

• Authentication: $ gcloud auth login rerun this in Google Cloud Notebook terminal when you are logged out and need the credential again.

Installation

Install the tensorflow_datasets to prepare sample dataset, and the grpcio-tools for querying against the index.

In [ ]:

! pip install -U grpcio-tools --user
! pip install -U tensorflow==2.9.1 --user
! pip install -U tensorflow-datasets --user

Download and install the latest (preview) version of the Vertex SDK for Python.

In [ ]:

! pip install -U git+https://github.com/googleapis/python-aiplatform.git@main-test --user

Install Cohere

In [ ]:

! pip install -U cohere --user

Restart the kernel (Colab)

After you install the additional packages, you need to restart the notebook kernel so it can find the packages.

In [ ]:

# Automatically restart kernel after installs
import os

if not os.getenv("IS_TESTING"):
    # Automatically restart kernel after installs
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

Creating Embeddings with Cohere

In [ ]:

# Lets start by loading the dataset with tensorflow-datasets
import tensorflow_datasets as tfds

dataset = tfds.load("ag_news_subset", split="train", shuffle_files=True)

In [ ]:

# For speed and cost considerations, lets limit the dataset to 1000 examples
df = tfds.as_dataframe(dataset.take(1000), tfds.builder("ag_news_subset").info)
df["text"] = df["description"].apply(lambda x: x.decode())

In [ ]:

# Finally, lets import cohere and use co.embed to create representations for these 1000 articles
import cohere

co = cohere.Client(COHERE_API_KEY)

In [ ]:

# running each of the examples through the embedding endpoint
response = co.embed(model="small", texts=list(df["text"].values))

cohere_embeddings = response.embeddings

Set up your Google Cloud project

The following steps are required, regardless of your notebook environment.

1. Select or create a Google Cloud project.

2. Make sure that billing is enabled for your project.

3. Enable the Vertex AI API and Compute Engine API, and Service Networking API.

4. Enter your project ID in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook.

Note: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $ into these commands.

Set your project ID

If you don't know your project ID, you may be able to get your project ID using gcloud.

In [ ]:

import os

PROJECT_ID = GOOGLE_PROJECT_ID

# Get your Google Cloud project ID from gcloud
if not os.getenv("IS_TESTING"):
    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null
    PROJECT_ID = shell_output[0]
    print("Project ID: ", PROJECT_ID)

Otherwise, set your project ID here.

In [ ]:

if PROJECT_ID == "" or PROJECT_ID is None:
    PROJECT_ID = GOOGLE_PROJECT_ID  # @param {type:"string"}

Create a Cloud Storage bucket

The following steps are required, regardless of your notebook environment.

Set the name of your Cloud Storage bucket below. It must be unique across all Cloud Storage buckets.

You may also change the REGION variable, which is used for operations throughout the rest of this notebook. Make sure to choose a region where Vertex AI services are available. You may not use a Multi-Regional Storage bucket for training with Vertex AI.

In [ ]:

from datetime import datetime

TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")

if BUCKET_NAME == "" or BUCKET_NAME is None or BUCKET_NAME == "gs://[your-bucket-name]":
    BUCKET_NAME = "gs://" + PROJECT_ID + "aip-" + TIMESTAMP

Only if your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.

In [ ]:

! gsutil mb -l $REGION $BUCKET_NAME

Finally, validate access to your Cloud Storage bucket by examining its contents:

In [ ]:

# this will not return anything if the bucket is empty
! gsutil ls -al $BUCKET_NAME

Import libraries and define constants

Import the Vertex AI (unified) client library into your Python environment.

In [ ]:

import time

import grpc
from google.cloud import aiplatform_v1beta1
from google.protobuf import struct_pb2

In [ ]:

ENDPOINT = "{}-aiplatform.googleapis.com".format(REGION)


AUTH_TOKEN = !gcloud auth print-access-token
PROJECT_NUMBER = !gcloud projects list --filter="PROJECT_ID:'{PROJECT_ID}'" --format='value(PROJECT_NUMBER)'
PROJECT_NUMBER = PROJECT_NUMBER[0]

PARENT = "projects/{}/locations/{}".format(PROJECT_ID, REGION)

print("ENDPOINT: {}".format(ENDPOINT))
print("PROJECT_ID: {}".format(PROJECT_ID))
print("REGION: {}".format(REGION))

!gcloud config set project {PROJECT_ID}
!gcloud config set ai_platform/region {REGION}

Prepare the Embeddings

This will take the embeddings generated with Cohere and format them to work with Matching Engine

Save the data in JSONL format.

In [ ]:

# This converts the list of embeddings to the json format expected by Matching Engine

with open("cohere_embeddings.json", "w") as f:
    for i, e in enumerate(cohere_embeddings):
        f.write('{"id":"' + str(i) + '",')
        f.write('"embedding":' + str(e) + "}")
        f.write("\n")

Upload the data to GCS.

In [ ]:

# NOTE: Everything in this GCS DIR will be DELETED before uploading the data.
# A CommandException is expected if no data is present

! gsutil rm -rf {BUCKET_NAME}/*

In [ ]:

! gsutil cp cohere_embeddings.json {BUCKET_NAME}/cohere_embeddings.json

In [ ]:

! gsutil ls {BUCKET_NAME}

Create Indexes

Create ANN Index (for Production Usage)

In [ ]:

index_client = aiplatform_v1beta1.IndexServiceClient(
    client_options=dict(api_endpoint=ENDPOINT)
)

In [ ]:

# Cohere small model is 1024 dimensions, update the dimension size if another model is being used
DIMENSIONS = 1024
DISPLAY_NAME = "cohere_embeddings"
DISPLAY_NAME_BRUTE_FORCE = DISPLAY_NAME + "_brute_force"

Create the ANN index configuration:

Please read the documentation to understand the various configuration parameters that can be used to tune the index

In [ ]:

treeAhConfig = struct_pb2.Struct(
    fields={
        "leafNodeEmbeddingCount": struct_pb2.Value(number_value=500),
        "leafNodesToSearchPercent": struct_pb2.Value(number_value=7),
    }
)

algorithmConfig = struct_pb2.Struct(
    fields={"treeAhConfig": struct_pb2.Value(struct_value=treeAhConfig)}
)

config = struct_pb2.Struct(
    fields={
        "dimensions": struct_pb2.Value(number_value=DIMENSIONS),
        "approximateNeighborsCount": struct_pb2.Value(number_value=150),
        "distanceMeasureType": struct_pb2.Value(string_value="DOT_PRODUCT_DISTANCE"),
        "algorithmConfig": struct_pb2.Value(struct_value=algorithmConfig),
    }
)

metadata = struct_pb2.Struct(
    fields={
        "config": struct_pb2.Value(struct_value=config),
        "contentsDeltaUri": struct_pb2.Value(string_value=BUCKET_NAME),
    }
)

ann_index = {
    "display_name": DISPLAY_NAME,
    "description": "Glove 100 ANN index",
    "metadata": struct_pb2.Value(struct_value=metadata),
}

In [ ]:

ann_index = index_client.create_index(parent=PARENT, index=ann_index)

In [ ]:

# Poll the operation until it's done successfullly.
# This will take some time (~30 minutes)

while True:
    if ann_index.done():
        break
    print("Poll the operation to create index...")
    time.sleep(60)

In [ ]:

INDEX_RESOURCE_NAME = ann_index.result().name
INDEX_RESOURCE_NAME

Create Brute Force Index (for Ground Truth)

The brute force index uses a naive brute force method to find the nearest neighbors. This method is not fast or efficient. Hence brute force indices are not recommended for production usage. They are to be used to find the "ground truth" set of neighbors, so that the "ground truth" set can be used to measure recall of the indices being tuned for production usage. To ensure an apples to apples comparison, the distanceMeasureType and featureNormType, dimensions of the brute force index should match those of the production indices being tuned.

Create the brute force index configuration:

In [ ]:

algorithmConfig = struct_pb2.Struct(
    fields={"bruteForceConfig": struct_pb2.Value(struct_value=struct_pb2.Struct())}
)

config = struct_pb2.Struct(
    fields={
        "dimensions": struct_pb2.Value(number_value=DIMENSIONS),
        "approximateNeighborsCount": struct_pb2.Value(number_value=150),
        "distanceMeasureType": struct_pb2.Value(string_value="DOT_PRODUCT_DISTANCE"),
        "algorithmConfig": struct_pb2.Value(struct_value=algorithmConfig),
    }
)

metadata = struct_pb2.Struct(
    fields={
        "config": struct_pb2.Value(struct_value=config),
        "contentsDeltaUri": struct_pb2.Value(string_value=BUCKET_NAME),
    }
)

brute_force_index = {
    "display_name": DISPLAY_NAME_BRUTE_FORCE,
    "description": "Glove 100 index (brute force)",
    "metadata": struct_pb2.Value(struct_value=metadata),
}

In [ ]:

brute_force_index = index_client.create_index(parent=PARENT, index=brute_force_index)

In [ ]:

# Poll the operation until it's done successfullly.
# This will take ~45 min.

while True:
    if brute_force_index.done():
        break
    print("Poll the operation to create index...")
    time.sleep(60)

In [ ]:

INDEX_BRUTE_FORCE_RESOURCE_NAME = brute_force_index.result().name
INDEX_BRUTE_FORCE_RESOURCE_NAME

Create an IndexEndpoint with VPC Network

In [ ]:

index_endpoint_client = aiplatform_v1beta1.IndexEndpointServiceClient(
    client_options=dict(api_endpoint=ENDPOINT)
)

In [ ]:

VPC_NETWORK_NAME = "projects/{}/global/networks/{}".format(PROJECT_NUMBER, NETWORK_NAME)
VPC_NETWORK_NAME

In [ ]:

index_endpoint = {
    "display_name": "index_endpoint_for_demo",
    "network": VPC_NETWORK_NAME,
}

In [ ]:

r = index_endpoint_client.create_index_endpoint(
    parent=PARENT, index_endpoint=index_endpoint
)

In [ ]:

r.result()

In [ ]:

INDEX_ENDPOINT_NAME = r.result().name
INDEX_ENDPOINT_NAME

Deploy Indexes

Deploy ANN Index

In [ ]:

DEPLOYED_INDEX_ID = "cohere_embedding_deployed"

In [ ]:

deploy_ann_index = {
    "id": DEPLOYED_INDEX_ID,
    "display_name": DEPLOYED_INDEX_ID,
    "index": INDEX_RESOURCE_NAME,
}

In [ ]:

r = index_endpoint_client.deploy_index(
    index_endpoint=INDEX_ENDPOINT_NAME, deployed_index=deploy_ann_index
)

In [ ]:

# Poll the operation until it's done successfullly.
while True:
    if r.done():
        break
    print("Poll the operation to deploy index...")
    time.sleep(60)

In [ ]:

r.result()

Deploy Brute Force Index

In [ ]:

DEPLOYED_BRUTE_FORCE_INDEX_ID = "cohere_brute_force_deployed"

In [ ]:

deploy_brute_force_index = {
    "id": DEPLOYED_BRUTE_FORCE_INDEX_ID,
    "display_name": DEPLOYED_BRUTE_FORCE_INDEX_ID,
    "index": INDEX_BRUTE_FORCE_RESOURCE_NAME,
}

In [ ]:

r = index_endpoint_client.deploy_index(
    index_endpoint=INDEX_ENDPOINT_NAME, deployed_index=deploy_brute_force_index
)

In [ ]:

# Poll the operation until it's done successfullly.

while True:
    if r.done():
        break
    print("Poll the operation to deploy index...")
    time.sleep(60)

In [ ]:

r.result()

Create Online Queries

After you built your indexes, you may query against the deployed index through the online querying gRPC API (Match service) within the virtual machine instances from the same region (for example 'us-central1' in this tutorial).

The way a client uses this gRPC API is by folowing steps:

• Write match_service.proto locally

• Compile the protocal buffer (see below)

• Obtain the index endpoint

• Use a code-generated stub to make the call, passing the parameter values

In [ ]:

!git clone https://github.com/googleapis/googleapis.git

In [ ]:

%%writefile match_service.proto

syntax = "proto3";

package google.cloud.aiplatform.container.v1beta1;

import "google/rpc/status.proto";

// MatchService is a Google managed service for efficient vector similarity
// search at scale.
service MatchService {
  // Returns the nearest neighbors for the query. If it is a sharded
  // deployment, calls the other shards and aggregates the responses.
  rpc Match(MatchRequest) returns (MatchResponse) {}

  // Returns the nearest neighbors for batch queries. If it is a sharded
  // deployment, calls the other shards and aggregates the responses.
  rpc BatchMatch(BatchMatchRequest) returns (BatchMatchResponse) {}
}

// Parameters for a match query.
message MatchRequest {
  // The ID of the DeploydIndex that will serve the request.
  // This MatchRequest is sent to a specific IndexEndpoint of the Control API,
  // as per the IndexEndpoint.network. That IndexEndpoint also has
  // IndexEndpoint.deployed_indexes, and each such index has an
  // DeployedIndex.id field.
  // The value of the field below must equal one of the DeployedIndex.id
  // fields of the IndexEndpoint that is being called for this request.
  string deployed_index_id = 1;

  // The embedding values.
  repeated float float_val = 2;

  // The number of nearest neighbors to be retrieved from database for
  // each query. If not set, will use the default from
  // the service configuration.
  int32 num_neighbors = 3;

  // The list of restricts.
  repeated Namespace restricts = 4;

  // Crowding is a constraint on a neighbor list produced by nearest neighbor
  // search requiring that no more than some value k' of the k neighbors
  // returned have the same value of crowding_attribute.
  // It's used for improving result diversity.
  // This field is the maximum number of matches with the same crowding tag.
  int32 per_crowding_attribute_num_neighbors = 5;

  // The number of neighbors to find via approximate search before
  // exact reordering is performed. If not set, the default value from scam
  // config is used; if set, this value must be > 0.
  int32 approx_num_neighbors = 6;

  // The fraction of the number of leaves to search, set at query time allows
  // user to tune search performance. This value increase result in both search
  // accuracy and latency increase. The value should be between 0.0 and 1.0. If
  // not set or set to 0.0, query uses the default value specified in
  // NearestNeighborSearchConfig.TreeAHConfig.leaf_nodes_to_search_percent.
  int32 leaf_nodes_to_search_percent_override = 7;
}

// Response of a match query.
message MatchResponse {
  message Neighbor {
    // The ids of the matches.
    string id = 1;

    // The distances of the matches.
    double distance = 2;
  }
  // All its neighbors.
  repeated Neighbor neighbor = 1;
}

// Parameters for a batch match query.
message BatchMatchRequest {
  // Batched requests against one index.
  message BatchMatchRequestPerIndex {
    // The ID of the DeploydIndex that will serve the request.
    string deployed_index_id = 1;

    // The requests against the index identified by the above deployed_index_id.
    repeated MatchRequest requests = 2;

    // Selects the optimal batch size to use for low-level batching. Queries
    // within each low level batch are executed sequentially while low level
    // batches are executed in parallel.
    // This field is optional, defaults to 0 if not set. A non-positive number
    // disables low level batching, i.e. all queries are executed sequentially.
    int32 low_level_batch_size = 3;
  }

  // The batch requests grouped by indexes.
  repeated BatchMatchRequestPerIndex requests = 1;
}

// Response of a batch match query.
message BatchMatchResponse {
  // Batched responses for one index.
  message BatchMatchResponsePerIndex {
    // The ID of the DeployedIndex that produced the responses.
    string deployed_index_id = 1;

    // The match responses produced by the index identified by the above
    // deployed_index_id. This field is set only when the query against that
    // index succeed.
    repeated MatchResponse responses = 2;

    // The status of response for the batch query identified by the above
    // deployed_index_id.
    google.rpc.Status status = 3;
  }

  // The batched responses grouped by indexes.
  repeated BatchMatchResponsePerIndex responses = 1;
}

// Namespace specifies the rules for determining the datapoints that are
// eligible for each matching query, overall query is an AND across namespaces.
message Namespace {
  // The string name of the namespace that this proto is specifying,
  // such as "color", "shape", "geo", or "tags".
  string name = 1;

  // The allowed tokens in the namespace.
  repeated string allow_tokens = 2;

  // The denied tokens in the namespace.
  // The denied tokens have exactly the same format as the token fields, but
  // represents a negation. When a token is denied, then matches will be
  // excluded whenever the other datapoint has that token.
  //
  // For example, if a query specifies {color: red, blue, !purple}, then that
  // query will match datapoints that are red or blue, but if those points are
  // also purple, then they will be excluded even if they are red/blue.
  repeated string deny_tokens = 3;
}

Compile the protocol buffer, and then match_service_pb2.py and match_service_pb2_grpc.py are generated.

In [ ]:

! python -m grpc_tools.protoc -I=. --proto_path=googleapis --python_out=. --grpc_python_out=. match_service.proto

Obtain the Private Endpoint:

In [ ]:

DEPLOYED_INDEX_SERVER_IP = (
    list(index_endpoint_client.list_index_endpoints(parent=PARENT))[0]
    .deployed_indexes[0]
    .private_endpoints.match_grpc_address
)
DEPLOYED_INDEX_SERVER_IP

Test your query:

In [ ]:

import match_service_pb2
import match_service_pb2_grpc

channel = grpc.insecure_channel("{}:10000".format(DEPLOYED_INDEX_SERVER_IP))
stub = match_service_pb2_grpc.MatchServiceStub(channel)

Test the search with a query embedded with Cohere

In [ ]:

raw_query = "Articles about the climate"

In [ ]:

query = co.embed(model="small", texts=[raw_query]).embeddings[0]

In [ ]:

# Test query
request = match_service_pb2.MatchRequest()
request.deployed_index_id = DEPLOYED_INDEX_ID
for val in query:
    request.float_val.append(val)

response = stub.Match(request)
response

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. You can also manually delete resources that you created by running the following code.

In [ ]:

index_client.delete_index(name=INDEX_RESOURCE_NAME)

In [ ]:

index_endpoint_client.delete_index_endpoint(name=INDEX_ENDPOINT_NAME)