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
15
warnings.filterwarnings("ignore")
16
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(
13
    texts=list(df["text"]),
14
    model=model_name,
15
    input_type=input_type_embed,
16
).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(
6
    example_id, 10, include_distances=True
7
)
8
9
# Format and print the text and distances
10
results = pd.DataFrame(
11
    data={
12
        "texts": df.iloc[similar_item_ids[0]]["text"],
13
        "distance": similar_item_ids[1],
14
    }
15
).drop(example_id)
16
17
# NOTE: Your results might look slightly different to ours.
18
print(f"Question:'{df.iloc[example_id]['text']}'\nNearest neighbors:")
19
print(results)

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(
6
    texts=[query], model=model_name, input_type=input_type_query
7
).embeddings
8
9
# Retrieve the nearest neighbors
10
similar_item_ids = search_index.get_nns_by_vector(
11
    query_embed[0], 10, include_distances=True
12
)
13
# Format the results
14
query_results = pd.DataFrame(
15
    data={
16
        "texts": df.iloc[similar_item_ids[0]]["text"],
17
        "distance": similar_item_ids[1],
18
    }
19
)
20
21
22
# NOTE: Your results might look slightly different to ours.
23
print(f"Query:'{query}'\nNearest neighbors:")
24
print(query_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
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 = (
15
    alt.Chart(df_explore)
16
    .mark_circle(size=60)
17
    .encode(
18
        x=alt.X("x", scale=alt.Scale(zero=False)),  #'x',
19
        y=alt.Y("y", scale=alt.Scale(zero=False)),
20
        tooltip=["text"],
21
    )
22
    .properties(width=700, height=400)
23
)
24
chart.interactive()