Semantic Search

1
# Install Cohere for embeddings, Umap to reduce embeddings to 2 dimensions, 
2
# Altair for visualization, Annoy for approximate nearest neighbor search
3
pip install cohere umap-learn altair annoy datasets tqdm

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
4
# Import into a pandas dataframe, take only the first 1000 rows
5
df = pd.DataFrame(dataset)[:1000]
6
7
# Preview the data to ensure it has loaded correctly
8
print(df.head(10))

1
# We'll set up the name of the model we want to use, the API key, and the input type.
2
# Create and retrieve a Cohere API key from dashboard.cohere.ai/welcome/register
3
# Paste your API key here. Remember to not share publicly
4
model_name = "embed-english-v3.0"
5
api_key = ""
6
input_type_embed = "search_document"
7
8
# Now we'll set up the cohere client.
9
co = cohere.Client(api_key)
10
11
# Get the embeddings
12
embeds = co.embed(texts=list(df['text']),
13
                  model=model_name,
14
                  input_type=input_type_embed).embeddings

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

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

1
query = "What is the tallest mountain in the world?"
2
input_type_query = "search_query"
3
4
# Get the query's embedding
5
query_embed = co.embed(texts=[query],
6
                  model=model_name,
7
                  input_type=input_type_query).embeddings
8
9
# Retrieve the nearest neighbors
10
similar_item_ids = search_index.get_nns_by_vector(query_embed[0],10,
11
                                                include_distances=True)
12
# Format the results
13
query_results = pd.DataFrame(data={'texts': df.iloc[similar_item_ids[0]]['text'], 
14
                             'distance': similar_item_ids[1]})
15
16
17
print(f"Query:'{query}'\nNearest neighbors:")
18
print(query_results) # NOTE: Your results might look slightly different to ours.

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
7
# Prepare the data to plot and interactive visualization
8
# using Altair
9
df_explore = pd.DataFrame(data={'text': df['text']})
10
df_explore['x'] = umap_embeds[:,0]
11
df_explore['y'] = umap_embeds[:,1]
12
13
# Plot
14
chart = alt.Chart(df_explore).mark_circle(size=60).encode(
15
    x=#'x',
16
    alt.X('x',
17
        scale=alt.Scale(zero=False)
18
    ),
19
    y=
20
    alt.Y('y',
21
        scale=alt.Scale(zero=False)
22
    ),
23
    tooltip=['text']
24
).properties(
25
    width=700,
26
    height=400
27
)
28
chart.interactive()