mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Andrew Ferlitsch
|
98d71c75f7 |
-373
@@ -1,373 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "KSP1duKDeaDR",
|
||||
"metadata": {
|
||||
"id": "KSP1duKDeaDR"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright 2023 Google LLC\n",
|
||||
"#\n",
|
||||
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
||||
"# you may not use this file except in compliance with the License.\n",
|
||||
"# You may obtain a copy of the License at\n",
|
||||
"#\n",
|
||||
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
||||
"#\n",
|
||||
"# Unless required by applicable law or agreed to in writing, software\n",
|
||||
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
||||
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
||||
"# See the License for the specific language governing permissions and\n",
|
||||
"# limitations under the License."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "67b2c5cc-8fc6-4082-9052-69fa0377d770",
|
||||
"metadata": {
|
||||
"id": "67b2c5cc-8fc6-4082-9052-69fa0377d770"
|
||||
},
|
||||
"source": [
|
||||
"# Semantic Search using Embeddings\n",
|
||||
"\n",
|
||||
"Semantic search is a type of search that uses the meaning of words and phrases to find relevant results.\n",
|
||||
"\n",
|
||||
"In this tutorial, we will demonstrate how to do semantic search with embeddings generated from the news text and using [Google ScaNN: Efficient Vector Similarity Search](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) to retrieve the most relevant news semantically.\n",
|
||||
"\n",
|
||||
"## Pre-requisites:\n",
|
||||
"- Vertex LLM SDK\n",
|
||||
"- ScaNN [github](https://github.com/google-research/google-research/tree/master/scann)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "FyyMdUeAJIVv",
|
||||
"metadata": {
|
||||
"id": "FyyMdUeAJIVv"
|
||||
},
|
||||
"source": [
|
||||
"## Install Vertex LLM SDK\n",
|
||||
"\n",
|
||||
"DISCLAIMER: Text Embedding API is now in Experimental Preview. This release focuses on validating model prototypes and these models are not guaranteed to be released. Use of Text Embedding API is governed by the Google Cloud Terms of Service, the Pre-GA Offerings Terms of the GCP Service Specific Terms. The Acceptance Use Policy, and the Generative AI Prohibited Use Policy. Vertex Text Embedding API’s features may be unstable, change in backward-incompatible ways, and are not guaranteed to be released. There are no SLAs provided and no technical support obligations. GCP’s Cloud Data Processing Addendum does not apply to Pre-GA Offerings and customers should not use Text Embedding API to process personal data or other data subject to legal or regulatory compliance requirements. See description of launch stage for details.\n",
|
||||
"\n",
|
||||
"The information in this documentation is provided to the customer on an “as is” and “with all faults” basis without any warranty of any kind, either express or implied. Google does not warrant or guarantee the correctness, accuracy or reliability of the information in here. In no event will Google or its affiliates or licensors be liable for any damage or harm to customers from customer’s use of these materials."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "h6KaE3XRJdxc",
|
||||
"metadata": {
|
||||
"id": "h6KaE3XRJdxc"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from google.colab import auth as google_auth\n",
|
||||
"\n",
|
||||
"google_auth.authenticate_user()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "snBUuUamoJPz",
|
||||
"metadata": {
|
||||
"id": "snBUuUamoJPz"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip3 install google-cloud-aiplatform>=1.25 \"shapely<2.0.0\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "zgVQcE0ewO8W",
|
||||
"metadata": {
|
||||
"id": "zgVQcE0ewO8W"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"PROJECT_ID = \"cloud-nl-llm-embedding\" # @param {type:\"string\"}\n",
|
||||
"LOCATION = \"us-central1\" # @param {type:\"string\"}\n",
|
||||
"\n",
|
||||
"import vertexai\n",
|
||||
"\n",
|
||||
"vertexai.init(project=PROJECT_ID, location=LOCATION)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4xFzXmPbY7FC",
|
||||
"metadata": {
|
||||
"id": "4xFzXmPbY7FC"
|
||||
},
|
||||
"source": [
|
||||
"**Attention**: you would need to restart runtime so that the right package is installed."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "xnfG88OPZI18",
|
||||
"metadata": {
|
||||
"id": "xnfG88OPZI18"
|
||||
},
|
||||
"source": [
|
||||
"## Import TextEmbeddingModel"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "1e6e1b98-a632-44a2-afb8-fc212018ef4f",
|
||||
"metadata": {
|
||||
"id": "1e6e1b98-a632-44a2-afb8-fc212018ef4f"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from vertexai.preview.language_models import TextEmbeddingModel\n",
|
||||
"\n",
|
||||
"model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1suA-1HuaGj6",
|
||||
"metadata": {
|
||||
"id": "1suA-1HuaGj6"
|
||||
},
|
||||
"source": [
|
||||
"## Install ScaNN Package"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "770255d3-54dd-48c8-bbdd-fbc0be41f085",
|
||||
"metadata": {
|
||||
"id": "770255d3-54dd-48c8-bbdd-fbc0be41f085"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install scann"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d",
|
||||
"metadata": {
|
||||
"id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d"
|
||||
},
|
||||
"source": [
|
||||
"## Imports packages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "412d00f1-08db-4880-8ced-52a9583757b8",
|
||||
"metadata": {
|
||||
"id": "412d00f1-08db-4880-8ced-52a9583757b8"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import json\n",
|
||||
"import time\n",
|
||||
"\n",
|
||||
"import numpy as np\n",
|
||||
"import pandas as pd\n",
|
||||
"import scann"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b",
|
||||
"metadata": {
|
||||
"id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b"
|
||||
},
|
||||
"source": [
|
||||
"## Create Embedding Dataset.\n",
|
||||
"\n",
|
||||
"The dataset is solely to demonstrate the use of the Text Embedding API with a vector database. It is not intended to be used for any other purpose, such as evaluating models. The dataset is small and does not represent a comprehensive sample of all possible text."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2OUg-Qf8iFta",
|
||||
"metadata": {
|
||||
"id": "2OUg-Qf8iFta"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!gsutil cp gs://cloud-samples-data/vertex-ai/dataset-management/datasets/bert_finetuning/wide_and_deep_trainer_container_tests_input.jsonl ."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "BNPapKXviHlE",
|
||||
"metadata": {
|
||||
"id": "BNPapKXviHlE"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"records = []\n",
|
||||
"with open(\"wide_and_deep_trainer_container_tests_input.jsonl\") as f:\n",
|
||||
" for line in f:\n",
|
||||
" record = json.loads(line)\n",
|
||||
" records.append(record)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "Z1Y9Bx2miJba",
|
||||
"metadata": {
|
||||
"id": "Z1Y9Bx2miJba"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Peek at the data.\n",
|
||||
"df = pd.DataFrame(records)\n",
|
||||
"df.head(50)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "1v7zUwoWiPl-",
|
||||
"metadata": {
|
||||
"id": "1v7zUwoWiPl-"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def get_embedding(text):\n",
|
||||
" get_embedding.counter += 1\n",
|
||||
" try:\n",
|
||||
" if get_embedding.counter % 100 == 0:\n",
|
||||
" time.sleep(3)\n",
|
||||
" return model.get_embeddings([text])[0].values\n",
|
||||
" except:\n",
|
||||
" return []\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"get_embedding.counter = 0\n",
|
||||
"\n",
|
||||
"# This may take several minutes to complete.\n",
|
||||
"df[\"embedding\"] = df[\"textContent\"].apply(lambda x: get_embedding(x))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ba4f49b6-65e1-49ea-988b-c4e195219deb",
|
||||
"metadata": {
|
||||
"id": "ba4f49b6-65e1-49ea-988b-c4e195219deb"
|
||||
},
|
||||
"source": [
|
||||
"## Create an Index"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "245bc8cd-038b-484a-acb2-3a705d4cc4cf",
|
||||
"metadata": {
|
||||
"id": "245bc8cd-038b-484a-acb2-3a705d4cc4cf"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"record_count = len(records)\n",
|
||||
"dataset = np.empty((record_count, 768))\n",
|
||||
"for i in range(record_count):\n",
|
||||
" dataset[i] = df.embedding[i]\n",
|
||||
"\n",
|
||||
"normalized_dataset = dataset / np.linalg.norm(dataset, axis=1)[:, np.newaxis]\n",
|
||||
"# configure ScaNN as a tree - asymmetric hash hybrid with reordering\n",
|
||||
"# anisotropic quantization as described in the paper; see README\n",
|
||||
"\n",
|
||||
"# use scann.scann_ops.build() to instead create a TensorFlow-compatible searcher\n",
|
||||
"searcher = (\n",
|
||||
" scann.scann_ops_pybind.builder(normalized_dataset, 10, \"dot_product\")\n",
|
||||
" .tree(\n",
|
||||
" num_leaves=record_count,\n",
|
||||
" num_leaves_to_search=record_count,\n",
|
||||
" training_sample_size=record_count,\n",
|
||||
" )\n",
|
||||
" .score_ah(2, anisotropic_quantization_threshold=0.2)\n",
|
||||
" .reorder(100)\n",
|
||||
" .build()\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9f1689eb-c27c-4566-afc7-7fbc55552aad",
|
||||
"metadata": {
|
||||
"id": "9f1689eb-c27c-4566-afc7-7fbc55552aad"
|
||||
},
|
||||
"source": [
|
||||
"## Queries the Index"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "494079df-d8f6-4a6a-b26b-0b0477791adc",
|
||||
"metadata": {
|
||||
"id": "494079df-d8f6-4a6a-b26b-0b0477791adc"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def search(query):\n",
|
||||
" start = time.time()\n",
|
||||
" query = model.get_embeddings([query])[0].values\n",
|
||||
" neighbors, distances = searcher.search(query, final_num_neighbors=3)\n",
|
||||
" end = time.time()\n",
|
||||
"\n",
|
||||
" for id, dist in zip(neighbors, distances):\n",
|
||||
" print(f\"[docid:{id}] [{dist}] -- {df.textContent[int(id)][:125]}...\")\n",
|
||||
" print(\"Latency (ms):\", 1000 * (end - start))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "50497681-6112-4147-a13e-afedb72b54f5",
|
||||
"metadata": {
|
||||
"id": "50497681-6112-4147-a13e-afedb72b54f5"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"search(\"tell me about shark or animal\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "c814f835-3e2a-4366-a334-ccc636f00b83",
|
||||
"metadata": {
|
||||
"id": "c814f835-3e2a-4366-a334-ccc636f00b83"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"search(\"tell me about an important moment or event in your life\")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "text_embedding_api_semantic_search_with_scann.ipynb",
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
Reference in New Issue
Block a user