1079 lines
36 KiB
Text
1079 lines
36 KiB
Text
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cb1537e6",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Vector Database Introduction\n",
|
|
"\n",
|
|
"This notebook takes you through a simple flow to download some data, embed it, and then index and search it using a selection of vector databases. This is a common requirement for customers who want to store and search our embeddings with their own data in a secure environment to support production use cases such as chatbots, topic modelling and more.\n",
|
|
"\n",
|
|
"### What is a Vector Database\n",
|
|
"\n",
|
|
"A vector database is a database made to store, manage and search embedding vectors. The use of embeddings to encode unstructured data (text, audio, video and more) as vectors for consumption by machine-learning models has exploded in recent years, due to the increasing effectiveness of AI in solving use cases involving natural language, image recognition and other unstructured forms of data. Vector databases have emerged as an effective solution for enterprises to deliver and scale these use cases.\n",
|
|
"\n",
|
|
"### Why use a Vector Database\n",
|
|
"\n",
|
|
"Vector databases enable enterprises to take many of the embeddings use cases we've shared in this repo (question and answering, chatbot and recommendation services, for example), and make use of them in a secure, scalable environment. Many of our customers make embeddings solve their problems at small scale but performance and security hold them back from going into production - we see vector databases as a key component in solving that, and in this guide we'll walk through the basics of embedding text data, storing it in a vector database and using it for semantic search.\n",
|
|
"\n",
|
|
"\n",
|
|
"### Demo Flow\n",
|
|
"The demo flow is:\n",
|
|
"- **Setup**: Import packages and set any required variables\n",
|
|
"- **Load data**: Load a dataset and embed it using OpenAI embeddings\n",
|
|
"- **Pinecone**\n",
|
|
" - *Setup*: Here we setup the Python client for Pinecone. For more details go [here](https://docs.pinecone.io/docs/quickstart)\n",
|
|
" - *Index Data*: We'll create an index with namespaces for __titles__ and __content__\n",
|
|
" - *Search Data*: We'll test out both namespaces with search queries to confirm it works\n",
|
|
"- **Weaviate**\n",
|
|
" - *Setup*: Here we setup the Python client for Weaviate. For more details go [here](https://weaviate.io/developers/weaviate/current/client-libraries/python.html)\n",
|
|
" - *Index Data*: We'll create an index with __title__ search vectors in it\n",
|
|
" - *Search Data*: We'll run a few searches to confirm it works\n",
|
|
"\n",
|
|
"Once you've run through this notebook you should have a basic understanding of how to setup and use vector databases, and can move on to more complex use cases making use of our embeddings."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e2b59250",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup\n",
|
|
"\n",
|
|
"Import the required libraries and set the embedding model that we'd like to use."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "5be94df6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import openai\n",
|
|
"\n",
|
|
"import tiktoken\n",
|
|
"from tenacity import retry, wait_random_exponential, stop_after_attempt\n",
|
|
"from typing import List, Iterator\n",
|
|
"import concurrent\n",
|
|
"from tqdm import tqdm\n",
|
|
"import pandas as pd\n",
|
|
"from datasets import load_dataset\n",
|
|
"import numpy as np\n",
|
|
"import os\n",
|
|
"\n",
|
|
"# Pinecone's client library for Python\n",
|
|
"import pinecone\n",
|
|
"\n",
|
|
"# Weaviate's client library for Python\n",
|
|
"import weaviate\n",
|
|
"\n",
|
|
"# I've set this to our new embeddings model, this can be changed to the embedding model of your choice\n",
|
|
"EMBEDDING_MODEL = \"text-embedding-ada-002\"\n",
|
|
"\n",
|
|
"# Ignore unclosed SSL socket warnings - optional in case you get these errors\n",
|
|
"import warnings\n",
|
|
"\n",
|
|
"warnings.filterwarnings(action=\"ignore\", message=\"unclosed\", category=ResourceWarning)\n",
|
|
"warnings.filterwarnings(\"ignore\", category=DeprecationWarning) "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e5d9d2e1",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Load data\n",
|
|
"\n",
|
|
"In this section we'll source the data for this task, embed it and format it for insertion into a vector database"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"id": "bd99e08e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Simple function to take in a list of text objects and return them as a list of embeddings\n",
|
|
"def get_embeddings(input: List):\n",
|
|
" response = openai.Embedding.create(\n",
|
|
" input=input,\n",
|
|
" model=EMBEDDING_MODEL,\n",
|
|
" )[\"data\"]\n",
|
|
" return [data[\"embedding\"] for data in response]\n",
|
|
"\n",
|
|
"def batchify(iterable, n=1):\n",
|
|
" l = len(iterable)\n",
|
|
" for ndx in range(0, l, n):\n",
|
|
" yield iterable[ndx : min(ndx + n, l)]\n",
|
|
"\n",
|
|
"# Function for batching and parallel processing the embeddings\n",
|
|
"def embed_corpus(\n",
|
|
" corpus: List[str],\n",
|
|
" batch_size=64,\n",
|
|
" num_workers=8,\n",
|
|
" max_context_len=8191,\n",
|
|
"):\n",
|
|
"\n",
|
|
" # Encode the corpus, truncating to max_context_len\n",
|
|
" encoding = tiktoken.get_encoding(\"cl100k_base\")\n",
|
|
" encoded_corpus = [\n",
|
|
" encoded_article[:max_context_len] for encoded_article in encoding.encode_batch(corpus)\n",
|
|
" ]\n",
|
|
"\n",
|
|
" # Calculate corpus statistics: the number of inputs, the total number of tokens, and the estimated cost to embed\n",
|
|
" num_tokens = sum(len(article) for article in encoded_corpus)\n",
|
|
" cost_to_embed_tokens = num_tokens / 1_000 * 0.0004\n",
|
|
" print(\n",
|
|
" f\"num_articles={len(encoded_corpus)}, num_tokens={num_tokens}, est_embedding_cost={cost_to_embed_tokens:.2f} USD\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Embed the corpus\n",
|
|
" with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:\n",
|
|
" \n",
|
|
" futures = [\n",
|
|
" executor.submit(get_embeddings, text_batch)\n",
|
|
" for text_batch in batchify(encoded_corpus, batch_size)\n",
|
|
" ]\n",
|
|
"\n",
|
|
" with tqdm(total=len(encoded_corpus)) as pbar:\n",
|
|
" for _ in concurrent.futures.as_completed(futures):\n",
|
|
" pbar.update(batch_size)\n",
|
|
"\n",
|
|
" embeddings = []\n",
|
|
" for future in futures:\n",
|
|
" data = future.result()\n",
|
|
" embeddings.extend(data)\n",
|
|
"\n",
|
|
" return embeddings"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0c1c73cb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# We'll use the datasets library to pull the Simple Wikipedia dataset for embedding\n",
|
|
"dataset = list(load_dataset(\"wikipedia\", \"20220301.simple\")[\"train\"])\n",
|
|
"# Limited to 25k articles for demo purposes\n",
|
|
"dataset = dataset[:25_000] "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"id": "e6ee90ce",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"num_articles=25000, num_tokens=12896881, est_embedding_cost=5.16 USD\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"25024it [01:11, 348.92it/s] "
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"CPU times: user 15.8 s, sys: 1.96 s, total: 17.8 s\n",
|
|
"Wall time: 1min 14s\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"%%time\n",
|
|
"# Embed the article text\n",
|
|
"dataset_embeddings = embed_corpus([article[\"text\"] for article in dataset])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"id": "850c7215",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"num_articles=25000, num_tokens=88300, est_embedding_cost=0.04 USD\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"25024it [00:21, 1164.97it/s] \n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Embed the article titles separately\n",
|
|
"title_embeddings = embed_corpus([article[\"title\"] for article in dataset])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 17,
|
|
"id": "1410daaa",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<div>\n",
|
|
"<style scoped>\n",
|
|
" .dataframe tbody tr th:only-of-type {\n",
|
|
" vertical-align: middle;\n",
|
|
" }\n",
|
|
"\n",
|
|
" .dataframe tbody tr th {\n",
|
|
" vertical-align: top;\n",
|
|
" }\n",
|
|
"\n",
|
|
" .dataframe thead th {\n",
|
|
" text-align: right;\n",
|
|
" }\n",
|
|
"</style>\n",
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: right;\">\n",
|
|
" <th></th>\n",
|
|
" <th>id</th>\n",
|
|
" <th>url</th>\n",
|
|
" <th>title</th>\n",
|
|
" <th>text</th>\n",
|
|
" <th>title_vector</th>\n",
|
|
" <th>content_vector</th>\n",
|
|
" <th>vector_id</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <th>0</th>\n",
|
|
" <td>1</td>\n",
|
|
" <td>https://simple.wikipedia.org/wiki/April</td>\n",
|
|
" <td>April</td>\n",
|
|
" <td>April is the fourth month of the year in the J...</td>\n",
|
|
" <td>[0.0010547508718445897, -0.020757636055350304,...</td>\n",
|
|
" <td>[-0.011253940872848034, -0.013491976074874401,...</td>\n",
|
|
" <td>0</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>1</th>\n",
|
|
" <td>2</td>\n",
|
|
" <td>https://simple.wikipedia.org/wiki/August</td>\n",
|
|
" <td>August</td>\n",
|
|
" <td>August (Aug.) is the eighth month of the year ...</td>\n",
|
|
" <td>[0.0009623901569284499, 0.0008108559413813055,...</td>\n",
|
|
" <td>[0.0003609954728744924, 0.007262262050062418, ...</td>\n",
|
|
" <td>1</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>2</th>\n",
|
|
" <td>6</td>\n",
|
|
" <td>https://simple.wikipedia.org/wiki/Art</td>\n",
|
|
" <td>Art</td>\n",
|
|
" <td>Art is a creative activity that expresses imag...</td>\n",
|
|
" <td>[0.0033528385683894157, 0.006173426751047373, ...</td>\n",
|
|
" <td>[-0.004959689453244209, 0.015772193670272827, ...</td>\n",
|
|
" <td>2</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>3</th>\n",
|
|
" <td>8</td>\n",
|
|
" <td>https://simple.wikipedia.org/wiki/A</td>\n",
|
|
" <td>A</td>\n",
|
|
" <td>A or a is the first letter of the English alph...</td>\n",
|
|
" <td>[0.015449387952685356, -0.013746200129389763, ...</td>\n",
|
|
" <td>[0.024894846603274345, -0.022186409682035446, ...</td>\n",
|
|
" <td>3</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>4</th>\n",
|
|
" <td>9</td>\n",
|
|
" <td>https://simple.wikipedia.org/wiki/Air</td>\n",
|
|
" <td>Air</td>\n",
|
|
" <td>Air refers to the Earth's atmosphere. Air is a...</td>\n",
|
|
" <td>[0.0222249086946249, -0.020463958382606506, -0...</td>\n",
|
|
" <td>[0.021524671465158463, 0.018522677943110466, -...</td>\n",
|
|
" <td>4</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>\n",
|
|
"</div>"
|
|
],
|
|
"text/plain": [
|
|
" id url title \\\n",
|
|
"0 1 https://simple.wikipedia.org/wiki/April April \n",
|
|
"1 2 https://simple.wikipedia.org/wiki/August August \n",
|
|
"2 6 https://simple.wikipedia.org/wiki/Art Art \n",
|
|
"3 8 https://simple.wikipedia.org/wiki/A A \n",
|
|
"4 9 https://simple.wikipedia.org/wiki/Air Air \n",
|
|
"\n",
|
|
" text \\\n",
|
|
"0 April is the fourth month of the year in the J... \n",
|
|
"1 August (Aug.) is the eighth month of the year ... \n",
|
|
"2 Art is a creative activity that expresses imag... \n",
|
|
"3 A or a is the first letter of the English alph... \n",
|
|
"4 Air refers to the Earth's atmosphere. Air is a... \n",
|
|
"\n",
|
|
" title_vector \\\n",
|
|
"0 [0.0010547508718445897, -0.020757636055350304,... \n",
|
|
"1 [0.0009623901569284499, 0.0008108559413813055,... \n",
|
|
"2 [0.0033528385683894157, 0.006173426751047373, ... \n",
|
|
"3 [0.015449387952685356, -0.013746200129389763, ... \n",
|
|
"4 [0.0222249086946249, -0.020463958382606506, -0... \n",
|
|
"\n",
|
|
" content_vector vector_id \n",
|
|
"0 [-0.011253940872848034, -0.013491976074874401,... 0 \n",
|
|
"1 [0.0003609954728744924, 0.007262262050062418, ... 1 \n",
|
|
"2 [-0.004959689453244209, 0.015772193670272827, ... 2 \n",
|
|
"3 [0.024894846603274345, -0.022186409682035446, ... 3 \n",
|
|
"4 [0.021524671465158463, 0.018522677943110466, -... 4 "
|
|
]
|
|
},
|
|
"execution_count": 17,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# We then store the result in another dataframe, and prep the data for insertion into a vector DB\n",
|
|
"article_df = pd.DataFrame(dataset)\n",
|
|
"article_df['title_vector'] = title_embeddings\n",
|
|
"article_df['content_vector'] = dataset_embeddings\n",
|
|
"article_df['vector_id'] = article_df.index\n",
|
|
"article_df['vector_id'] = article_df['vector_id'].apply(str)\n",
|
|
"article_df.head()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ed32fc87",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Pinecone\n",
|
|
"\n",
|
|
"Now we'll look to index these embedded documents in a vector database and search them. The first option we'll look at is **Pinecone**, a managed vector database which offers a cloud-native option.\n",
|
|
"\n",
|
|
"Before you proceed with this step you'll need to navigate to [Pinecone](pinecone.io), sign up and then save your API key as an environment variable titled ```PINECONE_API_KEY```.\n",
|
|
"\n",
|
|
"For section we will:\n",
|
|
"- Create an index with multiple namespaces for article titles and content\n",
|
|
"- Store our data in the index with separate searchable \"namespaces\" for article **titles** and **content**\n",
|
|
"- Fire some similarity search queries to verify our setup is working"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 18,
|
|
"id": "92e6152a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"api_key = os.getenv(\"PINECONE_API_KEY\")\n",
|
|
"pinecone.init(api_key=api_key)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "63b28543",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Create Index\n",
|
|
"\n",
|
|
"First we need to create an index, which we'll call `wikipedia-articles`. Once we have an index, we can create multiple namespaces, which can make a single index searchable for various use cases. For more details, consult [Pinecone documentation](https://docs.pinecone.io/docs/namespaces#:~:text=Pinecone%20allows%20you%20to%20partition,different%20subsets%20of%20your%20index.).\n",
|
|
"\n",
|
|
"If you want to batch insert to your index in parallel to increase insertion speed then there is a great guide in the Pinecone documentation on [batch inserts in parallel](https://docs.pinecone.io/docs/insert-data#sending-upserts-in-parallel)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 19,
|
|
"id": "0a71c575",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Models a simple batch generator that make chunks out of an input DataFrame\n",
|
|
"class BatchGenerator:\n",
|
|
" \n",
|
|
" \n",
|
|
" def __init__(self, batch_size: int = 10) -> None:\n",
|
|
" self.batch_size = batch_size\n",
|
|
" \n",
|
|
" # Makes chunks out of an input DataFrame\n",
|
|
" def to_batches(self, df: pd.DataFrame) -> Iterator[pd.DataFrame]:\n",
|
|
" splits = self.splits_num(df.shape[0])\n",
|
|
" if splits <= 1:\n",
|
|
" yield df\n",
|
|
" else:\n",
|
|
" for chunk in np.array_split(df, splits):\n",
|
|
" yield chunk\n",
|
|
"\n",
|
|
" # Determines how many chunks DataFrame contains\n",
|
|
" def splits_num(self, elements: int) -> int:\n",
|
|
" return round(elements / self.batch_size)\n",
|
|
" \n",
|
|
" __call__ = to_batches\n",
|
|
"\n",
|
|
"df_batcher = BatchGenerator(300)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 20,
|
|
"id": "7ea9ad46",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"['wikipedia-articles']"
|
|
]
|
|
},
|
|
"execution_count": 20,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# Pick a name for the new index\n",
|
|
"index_name = 'wikipedia-articles'\n",
|
|
"\n",
|
|
"# Check whether the index with the same name already exists - if so, delete it\n",
|
|
"if index_name in pinecone.list_indexes():\n",
|
|
" pinecone.delete_index(index_name)\n",
|
|
" \n",
|
|
"# Creates new index\n",
|
|
"pinecone.create_index(name=index_name, dimension=len(article_df['content_vector'][0]))\n",
|
|
"index = pinecone.Index(index_name=index_name)\n",
|
|
"\n",
|
|
"# Confirm our index was created\n",
|
|
"pinecone.list_indexes()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 22,
|
|
"id": "5daeba00",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading vectors to content namespace..\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Upsert content vectors in content namespace\n",
|
|
"print(\"Uploading vectors to content namespace..\")\n",
|
|
"for batch_df in df_batcher(article_df):\n",
|
|
" index.upsert(vectors=zip(batch_df.vector_id, batch_df.content_vector), namespace='content')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 23,
|
|
"id": "5fc1b083",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading vectors to title namespace..\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Upsert title vectors in title namespace\n",
|
|
"print(\"Uploading vectors to title namespace..\")\n",
|
|
"for batch_df in df_batcher(article_df):\n",
|
|
" index.upsert(vectors=zip(batch_df.vector_id, batch_df.title_vector), namespace='title')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 24,
|
|
"id": "f90c7fba",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'dimension': 1536,\n",
|
|
" 'index_fullness': 0.1,\n",
|
|
" 'namespaces': {'content': {'vector_count': 25000},\n",
|
|
" 'title': {'vector_count': 25000}},\n",
|
|
" 'total_vector_count': 50000}"
|
|
]
|
|
},
|
|
"execution_count": 24,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# Check index size for each namespace to confirm all of our docs have loaded\n",
|
|
"index.describe_index_stats()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "2da40a69",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Search data\n",
|
|
"\n",
|
|
"Now we'll enter some dummy searches and check we get decent results back"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 25,
|
|
"id": "d701b3c7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# First we'll create dictionaries mapping vector IDs to their outputs so we can retrieve the text for our search results\n",
|
|
"titles_mapped = dict(zip(article_df.vector_id,article_df.title))\n",
|
|
"content_mapped = dict(zip(article_df.vector_id,article_df.text))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 28,
|
|
"id": "3c8c2aa1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def query_article(query, namespace, top_k=5):\n",
|
|
" '''Queries an article using its title in the specified\n",
|
|
" namespace and prints results.'''\n",
|
|
"\n",
|
|
" # Create vector embeddings based on the title column\n",
|
|
" embedded_query = openai.Embedding.create(\n",
|
|
" input=query,\n",
|
|
" model=EMBEDDING_MODEL,\n",
|
|
" )[\"data\"][0]['embedding']\n",
|
|
"\n",
|
|
" # Query namespace passed as parameter using title vector\n",
|
|
" query_result = index.query(embedded_query, \n",
|
|
" namespace=namespace, \n",
|
|
" top_k=top_k)\n",
|
|
"\n",
|
|
" # Print query results \n",
|
|
" print(f'\\nMost similar results to {query} in \"{namespace}\" namespace:\\n')\n",
|
|
" if not query_result.matches:\n",
|
|
" print('no query result')\n",
|
|
" \n",
|
|
" matches = query_result.matches\n",
|
|
" ids = [res.id for res in matches]\n",
|
|
" scores = [res.score for res in matches]\n",
|
|
" df = pd.DataFrame({'id':ids, \n",
|
|
" 'score':scores,\n",
|
|
" 'title': [titles_mapped[_id] for _id in ids],\n",
|
|
" 'content': [content_mapped[_id] for _id in ids],\n",
|
|
" })\n",
|
|
" \n",
|
|
" counter = 0\n",
|
|
" for k,v in df.iterrows():\n",
|
|
" counter += 1\n",
|
|
" print(f'{v.title} (score = {v.score})')\n",
|
|
" \n",
|
|
" print('\\n')\n",
|
|
"\n",
|
|
" return df"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 29,
|
|
"id": "67b3584d",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n",
|
|
"Most similar results to modern art in Europe in \"title\" namespace:\n",
|
|
"\n",
|
|
"Museum of Modern Art (score = 0.875286043)\n",
|
|
"Western Europe (score = 0.867383599)\n",
|
|
"Renaissance art (score = 0.864250064)\n",
|
|
"Pop art (score = 0.860506058)\n",
|
|
"Northern Europe (score = 0.854678154)\n",
|
|
"\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"query_output = query_article('modern art in Europe','title')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 30,
|
|
"id": "3e7ac79b",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n",
|
|
"Most similar results to Famous battles in Scottish history in \"content\" namespace:\n",
|
|
"\n",
|
|
"Battle of Bannockburn (score = 0.869324744)\n",
|
|
"Wars of Scottish Independence (score = 0.861479)\n",
|
|
"1651 (score = 0.852555931)\n",
|
|
"First War of Scottish Independence (score = 0.84969604)\n",
|
|
"Robert I of Scotland (score = 0.846192539)\n",
|
|
"\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"content_query_output = query_article(\"Famous battles in Scottish history\",'content')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d939342f",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Weaviate\n",
|
|
"\n",
|
|
"The other vector database option we'll explore here is **Weaviate**, which offers both a managed, SaaS option like Pinecone, as well as a self-hosted option. As we've already looked at a cloud vector database, we'll try the self-hosted option here.\n",
|
|
"\n",
|
|
"For this we will:\n",
|
|
"- Set up a local deployment of Weaviate\n",
|
|
"- Create indices in Weaviate\n",
|
|
"- Store our data there\n",
|
|
"- Fire some similarity search queries\n",
|
|
"- Try a real use case"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "bfdfe260",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Setup\n",
|
|
"\n",
|
|
"To get Weaviate running locally we used Docker and followed the instructions contained in this article: https://weaviate.io/developers/weaviate/current/installation/docker-compose.html\n",
|
|
"\n",
|
|
"For an example docker-compose.yaml file please refer to `./weaviate/docker-compose.yaml` in this repo\n",
|
|
"\n",
|
|
"You can start Weaviate up locally by navigating to this directory and running `docker-compose up -d `"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 33,
|
|
"id": "b9ea472d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"client = weaviate.Client(\"http://localhost:8080/\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 34,
|
|
"id": "13be220d",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'classes': []}"
|
|
]
|
|
},
|
|
"execution_count": 34,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"client.schema.delete_all()\n",
|
|
"client.schema.get()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 35,
|
|
"id": "73d33184",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"True"
|
|
]
|
|
},
|
|
"execution_count": 35,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"client.is_ready()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "03a926b9",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Index data\n",
|
|
"\n",
|
|
"In Weaviate you create __schemas__ to capture each of the entities you will be searching. \n",
|
|
"\n",
|
|
"In this case we'll create a schema called **Article** with the **title** vector from above included for us to search by.\n",
|
|
"\n",
|
|
"The next few steps closely follow the documents Weaviate provides [here](https://weaviate.io/developers/weaviate/current/tutorials/how-to-use-weaviate-without-modules.htm)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 36,
|
|
"id": "e868d143",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'classes': [{'class': 'Article',\n",
|
|
" 'invertedIndexConfig': {'bm25': {'b': 0.75, 'k1': 1.2},\n",
|
|
" 'cleanupIntervalSeconds': 60,\n",
|
|
" 'stopwords': {'additions': None, 'preset': 'en', 'removals': None}},\n",
|
|
" 'properties': [{'dataType': ['text'],\n",
|
|
" 'description': 'Title of the article',\n",
|
|
" 'name': 'title',\n",
|
|
" 'tokenization': 'word'},\n",
|
|
" {'dataType': ['text'],\n",
|
|
" 'description': 'Contents of the article',\n",
|
|
" 'name': 'content',\n",
|
|
" 'tokenization': 'word'}],\n",
|
|
" 'shardingConfig': {'virtualPerPhysical': 128,\n",
|
|
" 'desiredCount': 1,\n",
|
|
" 'actualCount': 1,\n",
|
|
" 'desiredVirtualCount': 128,\n",
|
|
" 'actualVirtualCount': 128,\n",
|
|
" 'key': '_id',\n",
|
|
" 'strategy': 'hash',\n",
|
|
" 'function': 'murmur3'},\n",
|
|
" 'vectorIndexConfig': {'skip': False,\n",
|
|
" 'cleanupIntervalSeconds': 300,\n",
|
|
" 'maxConnections': 64,\n",
|
|
" 'efConstruction': 128,\n",
|
|
" 'ef': -1,\n",
|
|
" 'dynamicEfMin': 100,\n",
|
|
" 'dynamicEfMax': 500,\n",
|
|
" 'dynamicEfFactor': 8,\n",
|
|
" 'vectorCacheMaxObjects': 2000000,\n",
|
|
" 'flatSearchCutoff': 40000,\n",
|
|
" 'distance': 'cosine'},\n",
|
|
" 'vectorIndexType': 'hnsw',\n",
|
|
" 'vectorizer': 'none'}]}"
|
|
]
|
|
},
|
|
"execution_count": 36,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"class_obj = {\n",
|
|
" \"class\": \"Article\",\n",
|
|
" \"vectorizer\": \"none\", # explicitly tell Weaviate not to vectorize anything, we are providing the vectors ourselves\n",
|
|
" \"properties\": [{\n",
|
|
" \"name\": \"title\",\n",
|
|
" \"description\": \"Title of the article\",\n",
|
|
" \"dataType\": [\"text\"]\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"name\": \"content\",\n",
|
|
" \"description\": \"Contents of the article\",\n",
|
|
" \"dataType\": [\"text\"]\n",
|
|
" }]\n",
|
|
"}\n",
|
|
"\n",
|
|
"# Create the schema in Weaviate\n",
|
|
"client.schema.create_class(class_obj)\n",
|
|
"\n",
|
|
"# Check that we've created it as intended\n",
|
|
"client.schema.get()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 39,
|
|
"id": "786d437f",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading vectors to article schema..\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Convert DF into a list of tuples\n",
|
|
"data_objects = []\n",
|
|
"for k,v in article_df.iterrows():\n",
|
|
" data_objects.append((v['title'],v['text'],v['title_vector'],v['vector_id']))\n",
|
|
"\n",
|
|
"# Upsert into article schema\n",
|
|
"print(\"Uploading vectors to article schema..\")\n",
|
|
"\n",
|
|
"# Store a list of UUIDs in case we want to use to refer back to the initial dataframe\n",
|
|
"uuids = []\n",
|
|
"\n",
|
|
"# Reuse our batcher from the Pinecone ingestion\n",
|
|
"for batch_df in df_batcher(article_df):\n",
|
|
" for k,v in batch_df.iterrows():\n",
|
|
" #print(articles)\n",
|
|
" uuid = client.data_object.create(\n",
|
|
" {\n",
|
|
" \"title\": v['title'],\n",
|
|
" \"content\": v['text']\n",
|
|
" },\n",
|
|
" \"Article\",\n",
|
|
" vector=v['title_vector']\n",
|
|
" )\n",
|
|
" uuids.append(uuid)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 47,
|
|
"id": "3658693c",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Cave Story\n",
|
|
"is a freeware video game released in 2004 for PC. It was thought of and created over five years by Daisuke Amaya, known by his pseudonym, or art name, Pixel. The game is an action-adventure game, and is similar to the Castlevania and Metroid games. It was first made in Japanese, and was translated to English by the fan translating group, Aeon Genesis.\n",
|
|
"\n",
|
|
"References \n",
|
|
"\n",
|
|
"Notes\n",
|
|
"\n",
|
|
"2004 video games\n",
|
|
"Amiga games\n",
|
|
"Dreamcast games\n",
|
|
"Freeware games\n",
|
|
"Indie video games\n",
|
|
"Nintendo 3DS games\n",
|
|
"Nintendo Switch games\n",
|
|
"MacOS games\n",
|
|
"Platform games\n",
|
|
"Sega Genesis games\n",
|
|
"Video games developed in Japan\n",
|
|
"Wii games\n",
|
|
"Windows games\n"
|
|
]
|
|
},
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'Aggregate': {'Article': [{'meta': {'count': 25000}}]}}"
|
|
]
|
|
},
|
|
"execution_count": 47,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# Test our insert has worked by checking one object\n",
|
|
"print(client.data_object.get()['objects'][0]['properties']['title'])\n",
|
|
"print(client.data_object.get()['objects'][0]['properties']['content'])\n",
|
|
"\n",
|
|
"# Test that all data has loaded\n",
|
|
"result = client.query.aggregate(\"Article\") \\\n",
|
|
" .with_fields('meta { count }') \\\n",
|
|
" .do()\n",
|
|
"result['data']"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "46050ca9",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Search Data\n",
|
|
"\n",
|
|
"As above, we'll fire some queries at our new Index and get back results based on the closeness to our existing vectors"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 48,
|
|
"id": "5acd5437",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def query_weaviate(query, schema, top_k=20):\n",
|
|
"\n",
|
|
" # Creates embedding vector from user query\n",
|
|
" embedded_query = openai.Embedding.create(\n",
|
|
" input=query,\n",
|
|
" model=EMBEDDING_MODEL,\n",
|
|
" )[\"data\"][0]['embedding']\n",
|
|
" \n",
|
|
" near_vector = {\"vector\": embedded_query}\n",
|
|
"\n",
|
|
" # Queries input schema with vectorised user query\n",
|
|
" query_result = client.query.get(schema,[\"title\",\"content\", \"_additional {certainty}\"]) \\\n",
|
|
" .with_near_vector(near_vector) \\\n",
|
|
" .with_limit(top_k) \\\n",
|
|
" .do()\n",
|
|
" \n",
|
|
" return query_result"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 49,
|
|
"id": "15def653",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"1. Museum of Modern Art (Score: 0.938)\n",
|
|
"2. Western Europe (Score: 0.934)\n",
|
|
"3. Renaissance art (Score: 0.932)\n",
|
|
"4. Pop art (Score: 0.93)\n",
|
|
"5. Northern Europe (Score: 0.927)\n",
|
|
"6. Hellenistic art (Score: 0.926)\n",
|
|
"7. Modernist literature (Score: 0.924)\n",
|
|
"8. Art film (Score: 0.922)\n",
|
|
"9. Central Europe (Score: 0.921)\n",
|
|
"10. Art (Score: 0.921)\n",
|
|
"11. European (Score: 0.921)\n",
|
|
"12. Byzantine art (Score: 0.92)\n",
|
|
"13. Postmodernism (Score: 0.92)\n",
|
|
"14. Eastern Europe (Score: 0.92)\n",
|
|
"15. Cubism (Score: 0.92)\n",
|
|
"16. Europe (Score: 0.919)\n",
|
|
"17. Impressionism (Score: 0.919)\n",
|
|
"18. Bauhaus (Score: 0.919)\n",
|
|
"19. Surrealism (Score: 0.919)\n",
|
|
"20. Expressionism (Score: 0.918)\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"query_result = query_weaviate('modern art in Europe','Article')\n",
|
|
"counter = 0\n",
|
|
"for article in query_result['data']['Get']['Article']:\n",
|
|
" counter += 1\n",
|
|
" print(f\"{counter}. { article['title']} (Score: {round(article['_additional']['certainty'],3) })\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 50,
|
|
"id": "93c4a696",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"1. Historic Scotland (Score: 0.947)\n",
|
|
"2. First War of Scottish Independence (Score: 0.946)\n",
|
|
"3. Battle of Bannockburn (Score: 0.946)\n",
|
|
"4. Wars of Scottish Independence (Score: 0.944)\n",
|
|
"5. Second War of Scottish Independence (Score: 0.94)\n",
|
|
"6. List of Scottish monarchs (Score: 0.937)\n",
|
|
"7. Scottish Borders (Score: 0.932)\n",
|
|
"8. Braveheart (Score: 0.929)\n",
|
|
"9. John of Scotland (Score: 0.929)\n",
|
|
"10. Guardians of Scotland (Score: 0.926)\n",
|
|
"11. Holyrood Abbey (Score: 0.925)\n",
|
|
"12. Scottish (Score: 0.925)\n",
|
|
"13. Scots (Score: 0.925)\n",
|
|
"14. Robert I of Scotland (Score: 0.924)\n",
|
|
"15. Scottish people (Score: 0.924)\n",
|
|
"16. Alexander I of Scotland (Score: 0.924)\n",
|
|
"17. Edinburgh Castle (Score: 0.924)\n",
|
|
"18. Robert Burns (Score: 0.923)\n",
|
|
"19. Battle of Bosworth Field (Score: 0.922)\n",
|
|
"20. David II of Scotland (Score: 0.922)\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"query_result = query_weaviate('Famous battles in Scottish history','Article')\n",
|
|
"counter = 0\n",
|
|
"for article in query_result['data']['Get']['Article']:\n",
|
|
" counter += 1\n",
|
|
" print(f\"{counter}. {article['title']} (Score: {round(article['_additional']['certainty'],3) })\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ad74202e",
|
|
"metadata": {},
|
|
"source": [
|
|
"Thanks for following along, you're now equipped to set up your own vector databases and use embeddings to do all kinds of cool things - enjoy! For more complex use cases please continue to work through other cookbook examples in this repo."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"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.10.9"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|