Basic Semantic Search with Cohere Models

1
#!pip install --upgrade cohere

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
model_name = "embed-v4.0"
2
api_key = ""
3
input_type_embed = "search_document"
4
5
co = cohere.Client(api_key)

1
dataset = load_dataset("trec", split="train")
2
3
df = pd.DataFrame(dataset)[:1000]
4
5
df.head(10)

1
embeds = co.embed(texts=list(df['text']),
2
                  model=model_name,
3
                  input_type=input_type_embed).embeddings

1
embeds = np.array(embeds)
2
embeds.shape

1
search_index = AnnoyIndex(embeds.shape[1], 'angular')
2
for i in range(len(embeds)):
3
    search_index.add_item(i, embeds[i])
4
5
search_index.build(10) # 10 trees
6
search_index.save('test.ann')

1
example_id = 92
2
3
similar_item_ids = search_index.get_nns_by_item(example_id,10,
4
                                                include_distances=True)
5
results = pd.DataFrame(data={'texts': df.iloc[similar_item_ids[0]]['text'],
6
                             'distance': similar_item_ids[1]}).drop(example_id)
7
8
print(f"Question:'{df.iloc[example_id]['text']}'\nNearest neighbors:")
9
results

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

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