In this article, we will look at how to use SingleStoreDB to store and query the OpenAI Wikipedia vector database dataset.
SingleStoreDB has supported a range of vector functions for some time, and these functions are ideally suited for storing embeddings, doing semantic search and using the data to provide context to OpenAI as part of the prompt. With this mechanism, we will be able to add “short-term” memory to ChatGPT.
The notebook file used in this article is available on GitHub.
In several previous articles, we have used some of the vector capabilities built into SingleStoreDB:
- Quick Tip: SingleStoreDB’s EUCLIDEAN_DISTANCE and JSON_ARRAY_PACK Functions
- Using SingleStore, Spark and Alternating Least Squares (ALS) to Build a Movie Recommender System
In this article, we’ll test the `JSON_ARRAY_PACK` and `DOT_PRODUCT` vector functions with the OpenAI Wikipedia Vector Database dataset.
There is an OpenAI notebook available on GitHub under an MIT License that tests several vector database systems. The tests can be run using local clients or in the cloud. In this article, we’ll use Singlestore Helios.
Create a Singlestore Helios Account
A previous article showed the steps required to create a free Singlestore Helios account. We’ll use the following settings:
- Workspace Group Name: OpenAI Demo Group
- Cloud Provider: AWS
- Region: US East 1 (N. Virginia)
- Workspace Name: openai-demo
- Size: S-00
- Advanced Settings: MarTech Application deselected
From the left-navigation pane, we’ll select DEVELOP 〉SQL Editor to create a new database, as follows:
1CREATE DATABASE IF NOT EXISTS openai_demo;
Import Notebook
From the left-navigation pane, we’ll select DEVELOP 〉Notebooks. In the top right of the web page we’ll select New Notebook 〉Import From File, as shown in Figure 1.
We’ll locate the .ipynb file downloaded from GitHub and import the file. We also need to ensure that we select the Connection and Database using the drop-down menus just above the Notebook, as shown in Figure 2.
OpenAI API Key
Before running the notebook, we must create an account on the OpenAI website. This provides some free credits. Since we will use embeddings, the cost will be minimal. We’ll also need to create an OpenAI API Key. This can be created from USER 〉API keys in our OpenAI account.
Work Through the Notebook
Let’s now work through the notebook. We’ll adhere to the flow and structure of the OpenAI notebook, and use some small code sections directly from the notebook where required.
Setup
First, some libraries:
1!pip install openai --quiet2!pip install tabulate --quiet3!pip install wget --quiet
Next, some imports:
1import openai2 3import pandas as pd4import wget5from ast import literal_eval6from sqlalchemy import *
and then the embedding model:
1EMBEDDING_MODEL = "text-embedding-ada-002"
1embeddings_url =2'https://cdn.openai.com/API/examples/data/vector_database_wikipedia_ar3ticles_embedded.zip'4 5# The file is ~700 MB so this will take some time6wget.download(embeddings_url)
and unpack it:
1import zipfile2 3with4zipfile.ZipFile("vector_database_wikipedia_articles_embedded.zip",5"r") as zip_ref:6 zip_ref.extractall("data")
Next, we’ll load the file into a Pandas Dataframe:
1article_df = pd.read_csv(2 "data/vector_database_wikipedia_articles_embedded.csv"3)
and we’ll take a look at the first few lines, as follows:
1article_df.head()
The next operation from the OpenAI notebook can take a while:
1# Read vectors from strings back into a list2article_df['title_vector'] =3article_df.title_vector.apply(literal_eval)4article_df['content_vector'] =5article_df.content_vector.apply(literal_eval)6 7# Set vector_id to be a string8article_df['vector_id'] = article_df['vector_id'].apply(str)
and then next, we’ll look at the Dataframe info:
1article_df.info(show_counts=True)
The result should be as follows:
1<class 'pandas.core.frame.DataFrame'>2RangeIndex: 25000 entries, 0 to 249993Data columns (total 7 columns):4 # Column Non-Null Count Dtype5--- ------ -------------- -----6 0 id 25000 non-null int647 1 url 25000 non-null object8 2 title 25000 non-null object9 3 text 25000 non-null object10 4 title_vector 25000 non-null object11 5 content_vector 25000 non-null object12 6 vector_id 25000 non-null object13dtypes: int64(1), object(6)
1%%sql2 3USE openai_demo;4DROP TABLE IF EXISTS wikipedia;5CREATE TABLE IF NOT EXISTS wikipedia (6 id INT PRIMARY KEY,7 url VARCHAR(255),8 title VARCHAR(100),9 text TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci,10 title_vector BLOB,11 content_vector BLOB,12 vector_id INT13);
Notice that we can enter SQL statements directly into the notebook cell using the %%sql magic command.
Populate table
We can populate our database table, as follows:
1db_connection = create_engine(connection_url)2 3# Prepare the statement4stmt = """5 INSERT INTO wikipedia (6 id,7 url,8 title,9 text,10 title_vector,11 content_vector,12 vector_id13 )14 VALUES (15 %s,16 %s,17 %s,18 %s,19 JSON_ARRAY_PACK_F64(%s),20 JSON_ARRAY_PACK_F64(%s),21 %s22 )23"""24 25# Convert the DataFrame to a NumPy record array26record_arr = article_df.to_records(index=False)27# Set the batch size28batch_size = 100029 30# Iterate over the rows of the record array in batches31for i in range(0, len(record_arr), batch_size):32 batch = record_arr[i:i+batch_size]33 values = [(34 row[0],35 row[1],36 row[2],37 row[3],38 str(row[4]),39 str(row[5]),40 int(row[6])41 ) for row in batch]42 db_connection.execute(stmt, values)
Loading the data should only take a few minutes. We can use other data loading methods, like pipelines, for larger datasets.
Search data
First, we’ll declare the OPENAI_API_KEY, as follows:
1openai.api_key = "〈OpenAI API Key〉"
Replace 〈OpenAI API Key&〉 with your key.
We’ll now define a Python function that will allow us to use either of the two vector columns in the database:
1from typing import Tuple, List2 3def search_wikipedia(4 query: str,5 column1: str,6 column2: str,7 num_rows: int = 108) -> Tuple[List[str], List[float]]:9 """Searches Wikipedia for the given query and returns the top10`num_rows` results.11 12 Args:13 query: The query to search for.14 column1: The name of the column in the Wikipedia database to15return for each result.16 column2: The name of the column in the Wikipedia database to17use as the score for each result.18 num_rows: The number of results to return.19 20 Returns:21 A list of the top `num_rows` results.22 """23 24 # Get the embedding of the query25 query_embedding_response = openai.Embedding.create(26 model=EMBEDDING_MODEL,27 input=query,28 )29 query_embedding = query_embedding_response["data"][0]["embedding"]30 31 # Create the SQL statement32 stmt = f"""33 SELECT34 {column1},35 DOT_PRODUCT_F64(JSON_ARRAY_PACK_F64(%s), {column2}) AS36score37 FROM wikipedia38 ORDER BY score DESC39 LIMIT %s40 """.format(column1=column1, column2=column2)41 42 # Execute the SQL statement43 results = db_connection.execute(stmt, [str(query_embedding),44num_rows])45 46 values = []47 scores = []48 49 # Separate the results into two lists50 for row in results:51 values.append(row[0])52 scores.append(row[1])53 54 # Return the results55 return values, scores
We can test SingleStoreDB using the two examples in the OpenAI notebook. First, we’ll use title and title_vector:
1values1, scores1 = search_wikipedia(2 query = "modern art in Europe",3 column1 = "title",4 column2 = "title_vector",5 num_rows = 56)
We’ll format the results using the following:
1from tabulate import tabulate2 3# Combine the values and scores lists into a list of tuples4# Each tuple contains a value and its corresponding score5table_data1 = list(zip(values1, scores1))6 7# Add a rank column to the table data8table_data1 = [(i + 1,) + data for i, data in enumerate(table_data1)]9 10# Create the table11table1 = tabulate(table_data1, headers=["Rank", "Title", "Score"])12 13# Print the table14print(table1)
The output should be similar to the following:
1Rank Title Score2------ -------------------- --------3 1 Museum of Modern Art 0.8751244 2 Western Europe 0.8675545 3 Renaissance art 0.8642096 4 Pop art 0.8603837 5 Northern Europe 0.854793
Next, we’ll use text and content_vector:
1values2, scores2 = search_wikipedia(2 query = "Famous battles in Scottish history",3 column1 = "text",4 column2 = "content_vector",5 num_rows = 56)
We’ll format the results using the following:
1# Combine the values and scores lists into a list of tuples2# Each tuple contains a value and its corresponding score3table_data2 = list(zip([value[:50] for value in values2], scores2))4 5# Add a rank column to the table data6table_data2 = [(i + 1,) + data for i, data in enumerate(table_data2)]7 8# Create the table9table2 = tabulate(table_data2, headers=["Rank", "Text", "Score"])10 11# Print the table12print(table2)
The output should be similar to the following:
1Rank Text Score2------ -------------------------------------------------- --------3 1 The Battle of Bannockburn, fought on 23 and 24 Jun 0.8693384 2 The Wars of Scottish Independence were a series of 0.861485 3 Events 0.8525336 January 1 – Charles II crowned King of7 4 The First War of Scottish Independence lasted from 0.8496428 5 Robert I of Scotland (11 July 1274 – 7 June 1329) 0.846184
Summary
In this article, we’ve seen that SingleStoreDB can store vectors with ease — and that we can also store other data types in the same table, such as numeric and text. With its powerful SQL and multi-model support, SingleStoreDB provides a one-stop solution for modern applications bringing both technical and business benefits through a single product.
If you are interested in further reading, check out these SingleStore blog posts:
- Why Your Vector Database Should Not be a Vector Database
- AI-Powered Semantic Search in SingleStoreDB
- Using Vector Functions for Image Matching in SQL with SingleStoreDB
