Skip to content

Commit

Permalink
Add matching engine vectorstore (#3350)
Browse files Browse the repository at this point in the history 
Co-authored-by: Tom Piaggio <tomaspiaggio@google.com>
Co-authored-by: scafati98 <jupyter@matchingengine.us-central1-a.c.scafati-joonix.internal>
Co-authored-by: scafati98 <scafatieugenio@gmail.com>
Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
  • Loading branch information
@hwchase17 @scafati98 @dev2049
5 people authored May 31, 2023
1 parent 8bcaca4 commit 470b282
Show file tree
Hide file tree
Showing 2 changed files with 787 additions and 0 deletions.
346 changes: 346 additions & 0 deletions docs/modules/indexes/vectorstores/examples/matchingengine.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,346 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "655b8f55-2089-4733-8b09-35dea9580695",
"metadata": {},
"source": [
"# MatchingEngine\n",
"\n",
"This notebook shows how to use functionality related to the GCP Vertex AI `MatchingEngine` vector database.\n",
"\n",
"> Vertex AI [Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview) provides the industry's leading high-scale low latency vector database. These vector databases are commonly referred to as vector similarity-matching or an approximate nearest neighbor (ANN) service.\n",
"\n",
"**Note**: This module expects an endpoint and deployed index already created as the creation time takes close to one hour. To see how to create an index refer to the section [Create Index and deploy it to an Endpoint](#create-index-and-deploy-it-to-an-endpoint)"
]
},
{
"cell_type": "markdown",
"id": "a9971578-0ae9-4809-9e80-e5f9d3dcc98a",
"metadata": {},
"source": [
"## Create VectorStore from texts"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f7c96da4-8d97-4f69-8c13-d2fcafc03b05",
"metadata": {},
"outputs": [],
"source": [
"from langchain.vectorstores import MatchingEngine"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "58b70880-edd9-46f3-b769-f26c2bcc8395",
"metadata": {},
"outputs": [],
"source": [
"texts = ['The cat sat on', 'the mat.', 'I like to', 'eat pizza for', 'dinner.', 'The sun sets', 'in the west.']\n",
"\n",
"\n",
"vector_store = MatchingEngine.from_components(\n",
" texts=texts,\n",
" project_id=\"<my_project_id>\",\n",
" region=\"<my_region>\",\n",
" gcs_bucket_uri=\"<my_gcs_bucket>\",\n",
" index_id=\"<my_matching_engine_index_id>\",\n",
" endpoint_id=\"<my_matching_engine_endpoint_id>\"\n",
")\n",
"\n",
"vector_store.add_texts(texts=texts)\n",
"\n",
"vector_store.similarity_search(\"lunch\", k=2)"
]
},
{
"cell_type": "markdown",
"id": "0e76e05c-d4ef-49a1-b1b9-2ea989a0eda3",
"metadata": {
"tags": []
},
"source": [
"## Create Index and deploy it to an Endpoint"
]
},
{
"cell_type": "markdown",
"id": "61935a91-5efb-48af-bb40-ea1e83e24974",
"metadata": {},
"source": [
"### Imports, Constants and Configs"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "421b66c9-5b8f-4ef7-821e-12886a62b672",
"metadata": {},
"outputs": [],
"source": [
"# Installing dependencies.\n",
"!pip install tensorflow \\\n",
" google-cloud-aiplatform \\\n",
" tensorflow-hub \\\n",
" tensorflow-text "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4e9cc02-371e-40a1-bce9-37ac8efdf2cb",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import json\n",
"\n",
"from google.cloud import aiplatform\n",
"import tensorflow_hub as hub\n",
"import tensorflow_text"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "352a05df-6532-4aba-a36f-603327a5bc5b",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"PROJECT_ID = \"<my_project_id>\"\n",
"REGION = \"<my_region>\"\n",
"VPC_NETWORK = \"<my_vpc_network_name>\"\n",
"PEERING_RANGE_NAME = \"ann-langchain-me-range\" # Name for creating the VPC peering.\n",
"BUCKET_URI = \"gs://<bucket_uri>\"\n",
"# The number of dimensions for the tensorflow universal sentence encoder. \n",
"# If other embedder is used, the dimensions would probably need to change.\n",
"DIMENSIONS = 512\n",
"DISPLAY_NAME = \"index-test-name\"\n",
"EMBEDDING_DIR = f\"{BUCKET_URI}/banana\"\n",
"DEPLOYED_INDEX_ID = \"endpoint-test-name\"\n",
"\n",
"PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
"PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
"VPC_NETWORK_FULL = f\"projects/{PROJECT_NUMBER}/global/networks/{VPC_NETWORK}\"\n",
"\n",
"# Change this if you need the VPC to be created.\n",
"CREATE_VPC = False"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "076e7931-f83e-4597-8748-c8004fd8de96",
"metadata": {},
"outputs": [],
"source": [
"# Set the project id\n",
"! gcloud config set project {PROJECT_ID}"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4265081b-a5b7-491e-8ac5-1e26975b9974",
"metadata": {},
"outputs": [],
"source": [
"# Remove the if condition to run the encapsulated code\n",
"if CREATE_VPC:\n",
" # Create a VPC network\n",
" ! gcloud compute networks create {VPC_NETWORK} --bgp-routing-mode=regional --subnet-mode=auto --project={PROJECT_ID}\n",
"\n",
" # Add necessary firewall rules\n",
" ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-icmp --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow icmp\n",
"\n",
" ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-internal --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow all --source-ranges 10.128.0.0/9\n",
"\n",
" ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-rdp --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow tcp:3389\n",
"\n",
" ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-ssh --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow tcp:22\n",
"\n",
" # Reserve IP range\n",
" ! gcloud compute addresses create {PEERING_RANGE_NAME} --global --prefix-length=16 --network={VPC_NETWORK} --purpose=VPC_PEERING --project={PROJECT_ID} --description=\"peering range\"\n",
"\n",
" # Set up peering with service networking\n",
" # Your account must have the \"Compute Network Admin\" role to run the following.\n",
" ! gcloud services vpc-peerings connect --service=servicenetworking.googleapis.com --network={VPC_NETWORK} --ranges={PEERING_RANGE_NAME} --project={PROJECT_ID}"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9dfbb847-fc53-48c1-b0f2-00d1c4330b01",
"metadata": {},
"outputs": [],
"source": [
"# Creating bucket.\n",
"! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
]
},
{
"cell_type": "markdown",
"id": "f9698068-3d2f-471b-90c3-dae3e4ca6f63",
"metadata": {},
"source": [
"### Using Tensorflow Universal Sentence Encoder as an Embedder"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "144007e2-ddf8-43cd-ac45-848be0458ba9",
"metadata": {},
"outputs": [],
"source": [
"# Load the Universal Sentence Encoder module\n",
"module_url = \"https://tfhub.dev/google/universal-sentence-encoder-multilingual/3\"\n",
"model = hub.load(module_url)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "94a2bdcb-c7e3-4fb0-8c97-cc1f2263f06c",
"metadata": {},
"outputs": [],
"source": [
"# Generate embeddings for each word\n",
"embeddings = model(['banana'])"
]
},
{
"cell_type": "markdown",
"id": "5a4e6e99-5e42-4e55-90f6-c03aae4fbf14",
"metadata": {},
"source": [
"### Inserting a test embedding"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "024c78f3-4663-4d8f-9f3c-b7d82073ada4",
"metadata": {},
"outputs": [],
"source": [
"initial_config = {\"id\": \"banana_id\", \"embedding\": [float(x) for x in list(embeddings.numpy()[0])]}\n",
"\n",
"with open(\"data.json\", \"w\") as f:\n",
" json.dump(initial_config, f)\n",
"\n",
"!gsutil cp data.json {EMBEDDING_DIR}/file.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a11489f4-5904-4fc2-9178-f32c2df0406d",
"metadata": {},
"outputs": [],
"source": [
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
]
},
{
"cell_type": "markdown",
"id": "e3c6953b-11f6-4803-bf2d-36fa42abf3c7",
"metadata": {},
"source": [
"### Creating Index"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c31c3c56-bfe0-49ec-9901-cd146f592da7",
"metadata": {},
"outputs": [],
"source": [
"my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
" display_name=DISPLAY_NAME,\n",
" contents_delta_uri=EMBEDDING_DIR,\n",
" dimensions=DIMENSIONS,\n",
" approximate_neighbors_count=150,\n",
" distance_measure_type=\"DOT_PRODUCT_DISTANCE\"\n",
")"
]
},
{
"cell_type": "markdown",
"id": "50770669-edf6-4796-9563-d1ea59cfa8e8",
"metadata": {},
"source": [
"### Creating Endpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "20c93d1b-a7d5-47b0-9c95-1aec1c62e281",
"metadata": {},
"outputs": [],
"source": [
"my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
" display_name=f\"{DISPLAY_NAME}-endpoint\",\n",
" network=VPC_NETWORK_FULL,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "b52df797-28db-4b4a-b79c-e8a274293a6a",
"metadata": {},
"source": [
"### Deploy Index"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "019a7043-ad11-4a48-bec7-18928547b2ba",
"metadata": {},
"outputs": [],
"source": [
"my_index_endpoint = my_index_endpoint.deploy_index(\n",
" index=my_index, \n",
" deployed_index_id=DEPLOYED_INDEX_ID\n",
")\n",
"\n",
"my_index_endpoint.deployed_indexes"
]
}
],
"metadata": {
"environment": {
"kernel": "python3",
"name": "common-cpu.m107",
"type": "gcloud",
"uri": "gcr.io/deeplearning-platform-release/base-cpu:m107"
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.1"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Loading

0 comments on commit 470b282

Please sign in to comment.