Searching all of Wikipedia


SingleStore Notebooks

Searching all of Wikipedia


This tutorial is meant for Standard & Premium Workspaces. You can't run this with a Free Starter Workspace due to restrictions on Storage. Create a Workspace using +group in the left nav & select Standard for this notebook. Gallery notebooks tagged with "Starter" are suitable to run on a Free Starter Workspace

In this notebook, we've embarked on the task to implement semantic search and Retrieval-Augmented Generation (RAG) across Wikipedia's extensive database, using SingleStore's indexed ANN search capabilities.

We have focused on a subset of Wikipedia, scraping 1,800 video game articles to obtain real OpenAI embeddings for about 40,000 paragraphs, supplemented by 10 million mock vectors for a scaled-up simulation. We then stored these vectors in a SingleStore database, applying different indexing options to enhance search efficiency. We show remarkable improvements in query response times, dropping to sub 100 milliseconds with indexing.

We also integrated this system with a RAG-based chat, where you can ask and retrieve contextual answers based on the video game information in the database. Additionally, I've shared the technical details and resources, including our Python code and data links, in this GitHub repository and an AWS S3 bucket (S3 URI: s3://wikipedia-video-game-data/video-game-embeddings(1).csv).

Part 0: Creating your workspace and database

For the purposes of this demo, we are only generating 10 million vectors to search through. By our estimations, Wikipedia contains around 160 million paragraphs as a whole. A quick heuristic for your workspace sizing:

  • 160 million vectors can be handled by an S-32 Workspace

  • 20 million vectors can be handled by an S-4 Workspace

You can now extrapolate the workspace size you require from the number of vectors you want to generate!

In [1]:

drop database if exists video_game_wikipedia;
create database video_game_wikipedia;
use video_game_wikipedia;
-- create our table (note the usage of the vector data type)
create table vecs(
id bigint(20),
url text default null,
paragraph text default null,
v vector(1536) not null,
shard key(id),
key(id) using hash,
fulltext (paragraph)

Part 1: Generating the Mock Vectors

A vecs table is created with fields for ID, URL, paragraph text, and the embedding vector.

Then, we define several functions:

  • randbetween: generates random floats in a specified range

  • gen_vector: creates a vector of given length filled with random values

  • normalize: adjusts a vector to unit length

  • norm1536: normalizes a vector of dimension 1536

  • nrandv1536: generates a normalized, random vector of dimension 1536

Finally, we populate vecs with 10,000,000 rows of mock vectors.

In [2]:

-- generates random floats in a specified range
create or replace function randbetween(a float, b float) returns float
return (rand()*(b - a) + a);
end ;

In [3]:

-- set the sql_mode so that the following function returns the expected result
set sql_mode = pipes_as_concat;
-- creates a vector of given length filled with random values
create or replace function gen_vector(length int) returns text as
declare s text = "[";
if length < 2 then
raise user_exception("length too short: " || length);
end if;
for i in 1..length-1 loop
s = s || randbetween(-1,1) || "," ;
end loop;
s = s || randbetween(-1,1) || "]";
return s;

In [4]:

-- adjusts a vector to unit length
create or replace function normalize(v blob) returns blob as
squares blob = vector_mul(v,v);
length float = sqrt(vector_elements_sum(squares));
return scalar_vector_mul(1/length, v);

In [5]:

-- normalizes a vector of dimension 1536
create or replace function norm1536(v vector(1536)) returns vector(1536) as
return normalize(v) :> vector(1536);

In [6]:

-- generates a normalized, random vector of dimension 1536
create or replace function nrandv1536() returns vector(1536) as
return norm1536(gen_vector(1536));

In [7]:

-- procedure to populate `vecs` with `num_rows` vectors
create or replace procedure insert_vectors(num_rows bigint) as
declare c int;
select count(*) into c from vecs;
insert into vecs (id, v)
select id + (select max(id) from vecs), nrandv1536()
from vecs
where id <= 128 * 1024; /* chunk size 128K so we can see progress */
select count(*) into c from vecs;
if c >= num_rows then
end if;
end loop;

In [8]:

-- run the procedure to populate `vecs` with 10,000,000 vectors
-- this will take around 20 min
insert into vecs (id, v) values (1, nrandv1536());
call insert_vectors(10000000);

As a quick aside, if you want to generate the full 160 million vectors, you simply have to change the num_rows to 160,000,000: call insert_vectors(160000000);

Part 2: Getting the Wikipedia video game data

We will use a SingleStore pipeline named wiki_pipeline to import data from an S3 bucket into vecs. The pipeline is configured to load data from a CSV file located at s3://wikipedia-video-game-data/video-game-embeddings(1).csv. Since the S3 bucket is open, the credentials section is left empty.

In [9]:

-- since the bucket is open, you can leave the credentials clause as it is
create or replace pipeline `wiki_pipeline` as
load data S3 's3://wikipedia-video-game-data/video-game-embeddings(1).csv'
config '{"region":"us-west-1"}'
credentials '{"aws_access_key_id": "",
"aws_secret_access_key": ""}'
skip duplicate key errors
into table `vecs`
format csv
fields terminated by ','
enclosed by '"'
lines terminated by '\r\n';

In [10]:

-- start the pipeline!
start pipeline `wiki_pipeline`;

In [11]:

-- monitor the pipeline!
from information_schema.PIPELINES_BATCHES_SUMMARY
order by BATCH_ID;

Part 3: Building the vector indexes

Now, we have all the data in our table vecs. Let's go ahead and build our vector index. SingleStore gives us many options for our index with many tunable parameters. We will stick with the IVF indexes with default parameters.

In [12]:

alter table vecs add vector index auto (v) INDEX_OPTIONS '{"index_type":"AUTO"}';

In [13]:

alter table vecs add vector index ivf_flat (v) INDEX_OPTIONS '{"index_type":"IVF_FLAT"}';

In [14]:

alter table vecs add vector index ivf_pq (v) INDEX_OPTIONS '{"index_type":"IVF_PQ"}';

Part 4: Testing our indexes

Now that we have indexed our 10M vector dataset, let us now run some queries to test index performance!

We have chosen a test vector whose paragraph is about Nintendo's Rad Racer video game. We will compare the performance of an exact K-nearest neighbor search to the searches with our ANN indexes.

As we will see, we get an order of magnitude improvement when using an index in comparison to the exact KNN search!

In [15]:

set @qv = (select v from vecs where id = 1125899906845489);
-- NO INDEX: exact kNN search
select paragraph, v <*> @qv as sim
from vecs
order by sim use index () desc
limit 5;

In [16]:

set @qv = (select v from vecs where id = 1125899906845489);
-- AUTO index
select paragraph, v <*> @qv as sim
from vecs
order by sim use index (auto) desc
limit 5;

In [17]:

set @qv = (select v from vecs where id = 1125899906845489);
select paragraph, v <*> @qv as sim
from vecs
order by sim use index (ivf_flat) desc
limit 5;

In [18]:

set @qv = (select v from vecs where id = 1125899906845489);
select paragraph, v <*> @qv as sim
from vecs
order by sim use index (ivf_pq) desc
limit 5;

Part 5: Hybrid Search in SingleStore

Let us now see how we can implement a "hybrid search" in SingleStore! This is going to be a query that combines two powerful tools: a fulltext search and a semantic search!

In [19]:

-- set the query vector
set @v_mario = (select v
from vecs where url = ""
order by id
limit 1);
-- building the hybrid search
with fts as(
select id, paragraph, match (paragraph) against ('Mario Kart') as score
from vecs
where match (paragraph) against ('Mario Kart')
order by score desc
limit 200
vs as (
select id, paragraph, v <*> @v_mario as score
from vecs
order by score use index (auto) desc
limit 200
.3 * ifnull(fts.score, 0) + .7 * vs.score as hybrid_score,
vs.score as vec_score,
ifnull(fts.score, 0) as ft_score
from fts full outer join vs
on =
order by hybrid_score desc
limit 5;

Part 6: Chatting with the Video Game data!

  • search_wiki_page Function:

    • Conducts semantic search in a database.

    • Uses embeddings from get_embedding for query input.

    • Finds top 'k' paragraphs in vecs table, ranked by similarity to query.

    • Measures and prints search execution time.

  • ask_wiki_page Function:

    • Utilizes results from search_wiki_page for chatbot input.

    • Generates a query for an OpenAI GPT model-based chatbot.

In [20]:

!pip3 install openai --quiet

In [21]:

import sqlalchemy as sa
from openai import OpenAI
import getpass
import os
import time
import json

In [22]:

# OpenAI connection
OPENAI_API_KEY = os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')
client = OpenAI()
EMBEDDING_MODEL = 'text-embedding-ada-002'
GPT_MODEL = 'gpt-3.5-turbo'
# SingleStore connection
engine = sa.create_engine(connection_url)
connection = engine.connect()

In [23]:

def get_embedding(text, model=EMBEDDING_MODEL):
'''Generates the OpenAI embedding from an input `text`.'''
if isinstance(text, str):
response = client.embeddings.create(input=[text], model=model)
return json.dumps([0].embedding)

In [24]:

def search_wiki_page(query, limit=5):
'''Returns a df of the top k matches to the query ordered by similarity.'''
query_embedding_vec = get_embedding(query)
statement = sa.text(
f'''select paragraph, v <*> :query_embedding :> vector(1536) AS similarity
from vecs
order by similarity use index (auto) desc
limit :limit;'''
print("Searching for matches...")
start_time = time.time()
results = connection.execute(statement, {"query_embedding": query_embedding_vec, "limit": limit})
end_time = time.time()
execution_time = end_time - start_time
print(f"Search complete in {execution_time} seconds.")
results_as_dict = results.fetchall()
return results_as_dict

In [25]:

def ask_wiki_page(query, limit=5, temp=0.0):
'''Uses RAG to answer a question from the wiki page'''
results = search_wiki_page(query, limit)
print("Asking Chatbot...")
prompt = f'''Excerpt from the conversation history:
Question: {query}
Based on the conversation history, try to provide the most accurate answer to the question.
Consider the details mentioned in the conversation history to formulate a response that is as
helpful and precise as possible. please provide links to WIKIPEDIA ARTICLES TO LOOK AT FOR MORE INFORMATION.
response =
{"role": "system", "content": "You are a helpful assistant who is answering questions about an article."},
{"role": "user", "content": prompt}
response_message = response.choices[0].message.content
return response_message

In [26]:

query = input('Ask me a question about video games!')


About this Template

Simulate searching through 10 million paragraphs in Wikipedia for information about video games, and generate answers using RAG!




This Notebook has been released under the Apache 2.0 open source license.