Semantic Search with Cohere Embed Jobs

1
import time
2
import cohere
3
import hnswlib
4
co = cohere.Client('COHERE_API_KEY')

1
dataset_file_path = "data/embed_jobs_sample_data.jsonl" # Full path - https://raw.githubusercontent.com/cohere-ai/notebooks/main/notebooks/data/embed_jobs_sample_data.jsonl
2
3
ds=co.create_dataset(
4
	name='sample_file',
5
	data=open(dataset_file_path, 'rb'),
6
	keep_fields = ['id','wiki_id'],
7
	dataset_type="embed-input"
8
	)

uploading file, starting validation...
sample-file-hca4x0 was uploaded
...

1
print(ds.await_validation())

cohere.Dataset {
    id: sample-file-hca4x0
    name: sample_file
    dataset_type: embed-input
    validation_status: validated
    created_at: 2024-01-13 02:51:48.215973
    updated_at: 2024-01-13 02:51:48.215973
    download_urls: ['']
    validation_error: None
    validation_warnings: []
}

1
job = co.create_embed_job(
2
    dataset_id=ds.id,
3
    input_type='search_document' ,
4
    model='embed-english-v3.0',
5
    embeddings_types=['float'])
6
7
job.wait() # poll the server until the job is completed

...
...

1
print(job)

cohere.EmbedJob {
    job_id: 792bbc1a-561b-48c2-8a97-0c80c1914ea8
    status: complete
    created_at: 2024-01-13T02:53:31.879719Z
    input_dataset_id: sample-file-hca4x0
    output_urls: None
    model: embed-english-v3.0
    truncate: RIGHT
    percent_complete: 100
    output: cohere.Dataset {
    id: embeded-sample-file-drtjf9
    name: embeded-sample-file
    dataset_type: embed-result
    validation_status: validated
    created_at: 2024-01-13 02:53:33.569362
    updated_at: 2024-01-13 02:53:33.569362
    download_urls: ['']
    validation_error: None
    validation_warnings: []
}
}

1
embeddings_file_path = 'embed_jobs_output.csv'
2
output_dataset=co.get_dataset(job.output.id)
3
output_dataset.save(filepath=embeddings_file_path, format="csv")

1
embeddings=[]
2
texts=[]
3
for record in output_dataset:
4
  embeddings.append(record['embeddings']['float'])
5
  texts.append(record['text'])

1
index = hnswlib.Index(space='ip', dim=1024)
2
index.init_index(max_elements=len(embeddings), ef_construction=512, M=64)
3
index.add_items(embeddings,list(range(len(embeddings))))

1
query = "What was the first youtube video about?"
2
3
query_emb=co.embed(
4
    texts=[query], model="embed-english-v3.0", input_type="search_query"
5
        ).embeddings
6
7
doc_index = index.knn_query(query_emb, k=10)[0][0]
8
9
docs_to_rerank = []
10
for index in doc_index:
11
  docs_to_rerank.append(texts[index])
12
13
final_result = co.rerank(
14
    query= query,
15
    documents=docs_to_rerank,
16
    model="rerank-english-v2.0",
17
    top_n=3)

1
for idx, r in enumerate(final_result):
2
  print(f"Document Rank: {idx + 1}, Document Index: {r.index}")
3
  print(f"Document: {r.document['text']}")
4
  print(f"Relevance Score: {r.relevance_score:.5f}")
5
  print("\n")

Document Rank: 1, Document Index: 0
Document: YouTube began as a venture capital–funded technology startup. Between November 2005 and April 2006, the company raised money from various investors, with Sequoia Capital, $11.5 million, and Artis Capital Management, $8 million, being the largest two. YouTube's early headquarters were situated above a pizzeria and a Japanese restaurant in San Mateo, California. In February 2005, the company activated codice_1. The first video was uploaded April 23, 2005. Titled "Me at the zoo", it shows co-founder Jawed Karim at the San Diego Zoo and can still be viewed on the site. In May, the company launched a public beta and by November, a Nike ad featuring Ronaldinho became the first video to reach one million total views. The site launched officially on December 15, 2005, by which time the site was receiving 8 million views a day. Clips at the time were limited to 100 megabytes, as little as 30 seconds of footage.
Relevance Score: 0.94815
Document Rank: 2, Document Index: 1
Document: Karim said the inspiration for YouTube first came from the Super Bowl XXXVIII halftime show controversy when Janet Jackson's breast was briefly exposed by Justin Timberlake during the halftime show. Karim could not easily find video clips of the incident and the 2004 Indian Ocean Tsunami online, which led to the idea of a video-sharing site. Hurley and Chen said that the original idea for YouTube was a video version of an online dating service, and had been influenced by the website Hot or Not. They created posts on Craigslist asking attractive women to upload videos of themselves to YouTube in exchange for a $100 reward. Difficulty in finding enough dating videos led to a change of plans, with the site's founders deciding to accept uploads of any video.
Relevance Score: 0.91626
Document Rank: 3, Document Index: 2
Document: YouTube was not the first video-sharing site on the Internet; Vimeo was launched in November 2004, though that site remained a side project of its developers from CollegeHumor at the time and did not grow much, either. The week of YouTube's launch, NBC-Universal's "Saturday Night Live" ran a skit "Lazy Sunday" by The Lonely Island. Besides helping to bolster ratings and long-term viewership for "Saturday Night Live", "Lazy Sunday"'s status as an early viral video helped establish YouTube as an important website. Unofficial uploads of the skit to YouTube drew in more than five million collective views by February 2006 before they were removed when NBCUniversal requested it two months later based on copyright concerns. Despite eventually being taken down, these duplicate uploads of the skit helped popularize YouTube's reach and led to the upload of more third-party content. The site grew rapidly; in July 2006, the company announced that more than 65,000 new videos were being uploaded every day and that the site was receiving 100 million video views per day.
Relevance Score: 0.90665