A Deeper Dive Into Semantic Search

1
#title Import libraries (Run this cell to execute required code) {display-mode: "form"}
2
3
import cohere
4
import numpy as np
5
import re
6
import pandas as pd
7
from tqdm import tqdm
8
from datasets import load_dataset
9
import umap
10
import altair as alt
11
from sklearn.metrics.pairwise import cosine_similarity
12
from annoy import AnnoyIndex
13
import warnings
14
warnings.filterwarnings('ignore')
15
pd.set_option('display.max_colwidth', None)

1
# Get dataset
2
dataset = load_dataset("trec", split="train")
3
# Import into a pandas dataframe, take only the first 1000 rows
4
df = pd.DataFrame(dataset)[:1000]
5
# Preview the data to ensure it has loaded correctly
6
df.head(10)

1
# Paste your API key here. Remember to not share publicly
2
api_key = ''
3
4
# Create and retrieve a Cohere API key from dashboard.cohere.ai/welcome/register
5
co = cohere.Client(api_key)
6
7
# Get the embeddings
8
embeds = co.embed(texts=list(df['text']),
9
                  model='embed-english-v2.0').embeddings

1
# Create the search index, pass the size of embedding
2
search_index = AnnoyIndex(np.array(embeds).shape[1], 'angular')
3
# Add all the vectors to the search index
4
for i in range(len(embeds)):
5
    search_index.add_item(i, embeds[i])
6
search_index.build(10) # 10 trees
7
search_index.save('test.ann')

1
# Choose an example (we'll retrieve others similar to it)
2
example_id = 92
3
# Retrieve nearest neighbors
4
similar_item_ids = search_index.get_nns_by_item(example_id,10,
5
                                                include_distances=True)
6
# Format and print the text and distances
7
results = pd.DataFrame(data={'texts': df.iloc[similar_item_ids[0]]['text'],
8
                             'distance': similar_item_ids[1]}).drop(example_id)
9
print(f"Question:'{df.iloc[example_id]['text']}'\nNearest neighbors:")
10
results

1
# Output:
2
Question:'What are bear and bull markets ?'
3
Nearest neighbors:

1
query = "What is the tallest mountain in the world?"
2
3
# Get the query's embedding
4
query_embed = co.embed(texts=[query],
5
                  model="embed-english-v2.0").embeddings
6
7
# Retrieve the nearest neighbors
8
similar_item_ids = search_index.get_nns_by_vector(query_embed[0],10,
9
                                                include_distances=True)
10
# Format the results
11
results = pd.DataFrame(data={'texts': df.iloc[similar_item_ids[0]]['text'], 
12
                             'distance': similar_item_ids[1]})
13
14
15
print(f"Query:'{query}'\nNearest neighbors:")
16
results

1
#@title Plot the archive {display-mode: "form"}
2
3
# UMAP reduces the dimensions from 1024 to 2 dimensions that we can plot
4
reducer = umap.UMAP(n_neighbors=20) 
5
umap_embeds = reducer.fit_transform(embeds)
6
# Prepare the data to plot and interactive visualization
7
# using Altair
8
df_explore = pd.DataFrame(data={'text': df['text']})
9
df_explore['x'] = umap_embeds[:,0]
10
df_explore['y'] = umap_embeds[:,1]
11
12
# Plot
13
chart = alt.Chart(df_explore).mark_circle(size=60).encode(
14
    x=#'x',
15
    alt.X('x',
16
        scale=alt.Scale(zero=False)
17
    ),
18
    y=
19
    alt.Y('y',
20
        scale=alt.Scale(zero=False)
21
    ),
22
    tooltip=['text']
23
).properties(
24
    width=700,
25
    height=400
26
)
27
chart.interactive()