Doc2Vec Searching of Lang Database#

In this recipe, I will demonstrate the use of the gensim package and a simple doc2vec model trained on stories from the Lang coloured fairy books to support semantic retrieval of fairy stories.

The approach can be summarised as follows:

  • generate a vocabulary of terms representative of the search corpus;

  • generate a vector space where each dimension is a word in the vocabulary;

  • generate a vector for each document or search phrase;

  • retrieve documents based on similarity between document vector and search phrase vector.

The following recipe is inspired by How to make a search engine on Movies Description.

See also for ideas on a possible graphical user interface.

Connecting to the Database#

We’re going to work with our Lang fairy story database, so let’s set up a connection to it:

from sqlite_utils import Database

db_name = "demo.db"

db = Database(db_name)

Let’s remind ourselves of the database structure:

CREATE TABLE [books] (
   [book] TEXT,
   [title] TEXT,
   [text] TEXT,
   [last_para] TEXT,
   [first_line] TEXT,
   [provenance] TEXT,
   [chapter_order] INTEGER,
   PRIMARY KEY ([book], [title])
CREATE TABLE [books_metadata] (
   [title] TEXT,
   [year] INTEGER
    [title], [text],
CREATE TABLE 'books_fts_data'(id INTEGER PRIMARY KEY, block BLOB);
CREATE TABLE 'books_fts_idx'(segid, term, pgno, PRIMARY KEY(segid, term)) WITHOUT ROWID;
CREATE TABLE 'books_fts_docsize'(id INTEGER PRIMARY KEY, sz BLOB);
  INSERT INTO [books_fts] (rowid, [title], [text]) VALUES (new.rowid, new.[title], new.[text]);
  INSERT INTO [books_fts] ([books_fts], rowid, [title], [text]) VALUES('delete', old.rowid, old.[title], old.[text]);
  INSERT INTO [books_fts] ([books_fts], rowid, [title], [text]) VALUES('delete', old.rowid, old.[title], old.[text]);
  INSERT INTO [books_fts] (rowid, [title], [text]) VALUES (new.rowid, new.[title], new.[text]);

Recall that we can perform a full text search:

#q = 'king "three sons" gold'
q = 'hansel witch'
_q = f'SELECT title FROM books_fts WHERE books_fts MATCH {db.quote(q)} ;'

for row in db.query(_q):
Hansel And Grettel

We can randomly sample a selection of rows with a query of the following form:

# Via

_q_random_sample = """
WHERE rowid IN (SELECT rowid FROM books
                WHERE title NOT LIKE "Preface"
                ORDER BY random() LIMIT {});

for row in db.query(_q_random_sample.format(5)):
King Lindorm
Long, Broad, And Quickeye
The Princess Bella-Flor
The Hoodie-Crow.
Prince Featherhead And The Princess Celandine

Simple Model#

We could use an off-the-shelf model to process documents, or we can train our own model from our own documents so that the word vectors are aligned to our dataset. In a large corpus, we can train on a sample of documents if they are representative of the whole.

If we train against the whole dataset, we can search into the dataset directly from the model. If train the model on a partial collection, then we can only compare search phrases and documents that we have generated vectors for.

To create the model, it helps if we clean the documents, e.g. by decasing, and removing punctuation:

from gensim.parsing.preprocessing import preprocess_string
from gensim.parsing.preprocessing import strip_tags, strip_punctuation, strip_numeric, remove_stopwords

def clean_text(text):
    """Generate a cleaned, tokenised version of a text."""
    CUSTOM_FILTERS = [lambda x: x.lower(),
                      strip_tags, strip_punctuation,
                      strip_numeric, remove_stopwords]
    return preprocess_string(text, CUSTOM_FILTERS)

Apply the cleaning function to the text on the way in to creating the training corpus:

sample_corpus = db.query(_q_random_sample.format(9999))

sample_docs = [(clean_text(r['text']),
               f"{r['book']}::{r['title']}", #create a unique tag
               for r in sample_corpus]

# For the first doc, preview the first 5 cleaned words and title
sample_docs[0][0][:5], sample_docs[0][1], sample_docs[0][2]
(['time', 'certain', 'country', 'lived', 'king'],
 'The Blue Fairy Book::The Bronze Ring',
 'The Bronze Ring')
# The gensim model needs named tuples
# including at least a words and tags dimension
# Naively we can just use a document index count as the tag
from collections import namedtuple

StoryDoc = namedtuple('StoryDoc',
                      'words tags title')

sample_docs_training = []

for i, sample_doc in enumerate(sample_docs):
                                         [sample_doc[1]], # This must be a list
from gensim.models import Doc2Vec

# Define the parameters for building the model.
# We can also pass a list of documents
# via the first "documents" parameter
# and the model will be trained against those.
# Alternatively, create an empty model and train it later.
model = Doc2Vec(
                # dm: training algorithm;
                # 1: distributed memory/PV-DM;
                # 0: distributed bag of words (PV-DBOW)
                # vector_size: size of feature vectors
                # window: max dist between current & predicted word
                # hs: 1: hierarchical softmax;
                # hs: 0 : negative sampling if negative
                # min_count: ignore words w/ lower frequency
                # There is a risk to setting this too high
                # particularly if a search term is likely unique,
                # as it might be with a name. On the other hand,
                # for such situations, a simple search might be better?
                # sample: randomly downsample hi-frequency words
                # useful range: (0, 1e-5)

The model is built around a vocabulary extracted from the training document corpus.

# Build the model vocabulary

We can now train the model (this may take some time for a large corpus):

# It would be useful if we could display a progress bar for this
            epochs=100, start_alpha=0.01, end_alpha=0.01)

Rather than creating a model each time we want to use it, we can save the model and then load it as required:

# Save a model"lang_model.gensim")

# Load in a model
model = Doc2Vec.load("lang_model.gensim")

To retrieve a document matching a search phrase, we need to encode the search phrase and then try to find a matching document:

search_phrase = """
hansel and his sister were cast out by their wicked stepmother
and went into forest and met an evil witch

# Preprocess the search phrase
tokens = clean_text(search_phrase)


Generate a vector for the tokens:

# Generate the vector representation from the model
search_vector = model.infer_vector(tokens, alpha=0.001, epochs = 50)

We can now search for related documents from the original training set based on how well their vectors match the vector generated for the search phrase:

# Find the top 10 matches
matches = model.dv.most_similar([search_vector], topn=10)
# To rank every document from the training corpus
# set: topn=model.docvecs.count

# The response gives the original training document ids and match scores
[('The Blue Fairy Book::Hansel And Grettel', 0.728550136089325),
 ('The Yellow Fairy Book::The Witch', 0.5643637180328369),
 ('The Yellow Fairy Book::The White Duck', 0.5299764275550842),
 ('The Red Fairy Book::The Three Dwarfs', 0.5016872882843018),
 ('The Red Fairy Book::The Twelve Brothers', 0.4773011803627014),
 ('The Yellow Fairy Book::The Crow', 0.46543797850608826),
 ('The Yellow Fairy Book::The Dead Wife', 0.4470345079898834),
 ('The Red Fairy Book::Mother Holle', 0.4352571666240692),
 ('The Red Fairy Book::Brother And Sister', 0.43186280131340027),
 ('The Yellow Fairy Book::The Nixy', 0.4286130964756012)]

Let’s try another one:

search_phrase = """
a poor orphan girl lives with her wicked stepmother and sisters
but then her fairy godmother appears and she goes to a ball 
and leaves at midnight
but loses her slipper then finally marries the prince

# Preprocess the search phrase
tokens = clean_text(search_phrase)
search_vector = model.infer_vector(tokens, alpha=0.01, epochs = 50)
model.dv.most_similar([search_vector], topn=10)
[('The Blue Fairy Book::Cinderella, Or The Little Glass Slipper',
 ('The Brown Fairy Book::Which Was The Foolishest?', 0.4111166298389435),
 ('The Green Fairy Book::Spindle, Shuttle, And Needle', 0.3964698314666748),
 ('The Lilac Fairy Book::The Believing Husbands', 0.39020511507987976),
 ('The Lilac Fairy Book::The Brownie Of The Lake', 0.38550546765327454),
 ('The Brown Fairy Book::Habogi', 0.3834902048110962),
 ('The Blue Fairy Book::Little Red Riding Hood', 0.37533891201019287),
 ('The Lilac Fairy Book::The Hoodie-Crow.', 0.37399569153785706),
 ('The Violet Fairy Book::The Child Who Came From An Egg', 0.3730131685733795),
 ('The Violet Fairy Book::The Maiden With The Wooden Helmet',

Note that the result is stochastic (has a random element) in the way that the search vector is inferred: if you rerun the query, you will likely generate a different search vector. As a consequence, the search results returned are likely differ in their order and match scores each time the query is run.

Creating a Search Tool#

The next step is to register a custom SQLite function that will generate a vector for a search term and return matching records on that basis.

from vtfunc import TableFunction

class SQLite_doc2Vec_Model_Search(TableFunction):
    params = ['query', 'threshold']
    columns = ['book', 'title', 'score']
    name = 'doc2vec_model_search'
    model = Doc2Vec.load("lang_model.gensim")
    def initialize(self, query=None, num=None, threshold=None):
        tokens = clean_text(query)
        search_vector = model.infer_vector(tokens, alpha=0.01, epochs = 50)
        scores = model.dv.most_similar([search_vector],
        if threshold:
            scores = [(t, s) for (t, s) in scores if s >= threshold ]

        self._iter = iter(scores)
    def iterate(self, idx):
        (tag, score) = next(self._iter)
        items = tag.split("::")
        return (items[0], items[1], score,)

# And register the function

The query searches over the model and can take various forms:

  • doc2vec_model_search("search phrase")

  • doc2vec_model_search("search phrase", MIN_SCORE)

For example:

model_query = f"""
FROM doc2vec_model_search('''{search_phrase}''', 0.45 );

for i in db.execute(model_query):
('The Blue Fairy Book', 'Cinderella, Or The Little Glass Slipper', 0.48834240436553955)

Saving Model Vectors into the Database#

If we look at the object type of one of the model vectors, we see that it is a numpy.ndarray, which can be easily represented as a list:


We can store this data in the SQLite database as a BLOB. To simplify the process of converting the array into and out of the appropriate format for storage in the database compared to its use as a gensim vector, we can register a custom handler for the numpy.ndarray object:

# Via
# See also:
import sqlite3
import numpy as np
import io

def adapt_array(arr):
    """ (SoulNibbler)
    out = io.BytesIO(), arr)
    return sqlite3.Binary(

def convert_array(text):
    out = io.BytesIO(text)
    return np.load(out)

# Converts np.array to TEXT when inserting
sqlite3.register_adapter(np.ndarray, adapt_array)

# Converts TEXT to np.array when selecting
sqlite3.register_converter("array", convert_array)

Now we need to reset the database to a connection that supports the custom handler we have just registered:

# Reset the database connection to used the parsed datatype
db.conn = sqlite3.connect(db_name, detect_types=sqlite3.PARSE_DECLTYPES)

We can now create a table with a custom “array” datatype:

# Give ourselves a clean slate

# sqlite_utils doesn't appear to support custom types (yet?!)
# The following errors on the "array" datatype
    "book": str,
    "title": str,
    "tag": str, # a unique key derived from book and title
    "vector": "array",
}, pk=("book", "title"),
    # The following is not currently supported by sqlite_utils
   #foreign_keys=[ (("book", "title"), "books", ("book", "title"))] # local-table-id, foreign-table, foreign-table-id]

# so we can create a table the old fashioned way...
vector_table_create = """
CREATE TABLE story_vectors 
    (tag TEXT PRIMARY KEY, vector array, book TEXT, title TEXT );

cur = db.conn.cursor()
<sqlite3.Cursor at 0x12d1fbf10>

We can generate a list of dictionaries, one per record used to train the model, that can then be added directly to the story_vectors database table:

xrecords = []

for tag in model.dv.index_to_key:
    xrecords.append({'book': tag.split('::')[0],
                     'title': tag.split('::')[1],
                     'tag': tag, 'vector': model.dv[tag]})

And add the records directly to the database:

<Table story_vectors (tag, vector, book, title)>

Let’s pull an example record back showing just the first few elements of the vector associated with the record:

_q = f'SELECT * FROM story_vectors LIMIT 1;'

for row in db.query(_q):
    print(row['tag'], row['book'], row['title'], row['vector'][:10])
The Blue Fairy Book::The Bronze Ring The Blue Fairy Book The Bronze Ring [-1.1545537  -1.5133394  -0.9199181   3.0243099   0.9392002   2.3881683
 -1.5769789  -0.30107376 -0.3657083   1.5464693 ]
import pandas as pd

_q = f'SELECT * FROM story_vectors;'
df = pd.read_sql(_q, db.conn)

tag vector book title
0 The Blue Fairy Book::The Bronze Ring [-1.1545537, -1.5133394, -0.9199181, 3.0243099... The Blue Fairy Book The Bronze Ring
1 The Blue Fairy Book::Prince Hyacinth And The D... [-2.2750967, -0.05328803, -0.42831513, -0.4234... The Blue Fairy Book Prince Hyacinth And The Dear Little Princess
2 The Blue Fairy Book::East Of The Sun And West ... [-0.11260367, 1.1501226, 0.4695935, 0.6856275,... The Blue Fairy Book East Of The Sun And West Of The Moon

We can get the cosine similarity for each row relative to a search vector using the sklearn.metrics.pairwise.cosine_similarity applied to a dataframe of vectors we want to match against.

The cosine_similarity() function will happily accept two pandas dataframes, such as an N x M matrix of vectors we want to score against, and a 1 x M search vector matrics:

from sklearn.metrics.pairwise import cosine_similarity

# Grab the vectors and expand each on across columns
match_vectors_df = df['vector'].apply(pd.Series)
search_vector_df = pd.DataFrame(search_vector).T

df['score'] = cosine_similarity(match_vectors_df,

df[df['score']>0.45].sort_values("score", ascending=False).head()
tag vector book title score
6 The Blue Fairy Book::Cinderella, Or The Little... [-0.7598129, 1.4367288, 0.19395736, 0.573754, ... The Blue Fairy Book Cinderella, Or The Little Glass Slipper 0.486527

So it’s easy enough to create a custom function to search over the vectors table rather than the model:

class SQLite_doc2Vec_Table_Search(TableFunction):
    params = ['query', 'threshold']
    columns = ['book', 'title', 'score']
    name = 'doc2vec_search'
    # If we move this into the body, we can update the database
    # and respond to new rows added to story_vectors table
    _q = f'SELECT book, title, vector FROM story_vectors;'
    df = pd.read_sql(_q, db.conn)
    match_vectors_df = df['vector'].apply(pd.Series)
    def initialize(self, query=None, threshold=None):
        df = self.df
        tokens = clean_text(query)
        search_vector = model.infer_vector(tokens, alpha=0.01, epochs = 50)

        search_vector_df = pd.DataFrame(search_vector).T

        # Find cosine similarity
        df['score'] = cosine_similarity(self.match_vectors_df,
        # Apply minimum threshold if set
        _iterator = df[df['score']>=threshold] if threshold else df

        self._iter = _iterator[self.columns].itertuples(index=False,
    def iterate(self, idx):
        row = next(self._iter)
        return (row[0], row[1], row[2],)

# And register the function

Let’s try it out:

vector_query = f"""
FROM doc2vec_search("{search_phrase}")
WHERE score>0.4 ORDER BY score DESC LIMIT 3;

pd.read_sql(vector_query, db.conn)
book title score
0 The Blue Fairy Book Cinderella, Or The Little Glass Slipper 0.522628
1 The Red Fairy Book The Six Sillies 0.410686
2 The Brown Fairy Book Habogi 0.403348

Let’s try another one:

search_phrase2 = """
A girl lives with her grandmother and a wicked lord comes along and tells her
she can't get married unless she spins a funeral shroud and a wedding gown made from nettles.
As she spins, he gets ill, but can't die until she finishes it.

vector_query = f"""
FROM doc2vec_search("{search_phrase2}")
WHERE score>0.3 ORDER BY score DESC LIMIT 3;

pd.read_sql(vector_query, db.conn)
book title score
0 The Red Fairy Book The Nettle Spinner 0.639598
1 The Pink Fairy Book The Merry Wives 0.449067
2 The Blue Fairy Book Rumpelstiltzkin 0.444151
search_phrase3 = """
A girl in kitchen makes soup for king gets married.

vector_query = f"""
FROM doc2vec_search("{search_phrase3}")
WHERE score>0.3 ORDER BY score DESC LIMIT 10;

pd.read_sql(vector_query, db.conn)
book title score
0 The Yellow Fairy Book The Yellow Fairy Book 0.624495
1 The Blue Fairy Book Little Red Riding Hood 0.590124
2 The Red Fairy Book The Six Sillies 0.572264
3 The Blue Fairy Book Why The Sea Is Salt 0.547218
4 The Pink Fairy Book I Know What I Have Learned 0.535688
5 The Pink Fairy Book The Merry Wives 0.531312
6 The Red Fairy Book The True History Of Little Golden Hood 0.529552
7 The Pink Fairy Book Peter Bull 0.528966
8 The Pink Fairy Book The House In The Wood 0.517915
9 The Lilac Fairy Book A Lost Paradise 0.516366


To improve the doc2vec performance, it might be worth exploring a model that has a stricter minimum frequency for words in the corpus, but that also mixes a TF-IDF (term frequency, inverse document frequency) component in the ranking score?