Embeddings (Vectors, Search, Retrieval)

Semantic Search with Embeddings

This section provides examples on how to use the Embed endpoint to perform semantic search.

Semantic search solves the problem faced by the more traditional approach of lexical search, which is great at finding keyword matches, but struggles to capture the context or meaning of a piece of text.

PYTHON
import cohere
import numpy as np
co = cohere.ClientV2(
    api_key="YOUR_API_KEY"
)  # Get your free API key: https://dashboard.cohere.com/api-keys

The Embed endpoint takes in texts as input and returns embeddings as output.

For semantic search, there are two types of documents we need to turn into embeddings.

  • The list of documents to search from.
  • The query that will be used to search the documents.

Step 1: Embed the documents

We call the Embed endpoint using co.embed() and pass the required arguments:

  • texts: The list of texts
  • model: Here we choose embed-v4.0
  • input_type: We choose search_document to ensure the model treats these as the documents for search
  • embedding_types: We choose float to get a float array as the output

Step 2: Embed the query

Next, we add and embed a query. We choose search_query as the input_type to ensure the model treats this as the query (instead of documents) for search.

Step 3: Return the most similar documents

Next, we calculate and sort similarity scores between a query and document embeddings, then display the top N most similar documents. Here, we are using the numpy library for calculating similarity using a dot product approach.

### STEP 1: Embed the documents
# Define the documents
documents = [
    "Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.",
    "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee.",
    "Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!",
    "Working Hours Flexibility: We prioritize work-life balance. While our core hours are 9 AM to 5 PM, we offer flexibility to adjust as needed.",
]
# Constructing the embed_input object
embed_input = [
    {"content": [{"type": "text", "text": doc}]} for doc in documents
]
# Embed the documents
doc_emb = co.embed(
    inputs=embed_input,
    model="embed-v4.0",
    output_dimension=1024,
    input_type="search_document",
    embedding_types=["float"],
).embeddings.float
### STEP 2: Embed the query
# Add the user query
query = "How to connect with my teammates?"
query_input = [{"content": [{"type": "text", "text": query}]}]
# Embed the query
query_emb = co.embed(
    inputs=query_input,
    model="embed-v4.0",
    input_type="search_query",
    output_dimension=1024,
    embedding_types=["float"],
).embeddings.float
### STEP 3: Return the most similar documents
# Calculate similarity scores
scores = np.dot(query_emb, np.transpose(doc_emb))[0]
# Sort and filter documents based on scores
top_n = 2
top_doc_idxs = np.argsort(-scores)[:top_n]
# Display search results
for idx, docs_idx in enumerate(top_doc_idxs):
    print(f"Rank: {idx+1}")
    print(f"Document: {documents[docs_idx]}\n")

Here’s an example output:

Rank: 1
Document: Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!
Rank: 2
Document: Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged.

Content quality measure with Embed v4

A standard text embeddings model is optimized for only topic similarity between a query and candidate documents. But in many real-world applications, you have redundant information with varying content quality.

For instance, consider a user query of “COVID-19 Symptoms” and compare that to candidate document, “COVID-19 has many symptoms”. This document does not offer high-quality and rich information. However, with a typical embedding model, it will appear high on search results because it is highly similar to the query.

The Embed v4 model is trained to capture both content quality and topic similarity. Through this approach, a search system can extract richer information from documents and is robust against noise.

As an example below, give a query (“COVID-19 Symptoms”), the document with the highest quality (“COVID-19 symptoms can include: a high temperature or shivering…”) is ranked first.

Another document (“COVID-19 has many symptoms”) is arguably more similar to the query based on what information it contains, yet it is ranked lower as it doesn’t contain that much information.

This demonstrates how Embed v4 helps to surface high-quality documents for a given query.

PYTHON
### STEP 1: Embed the documents
documents = [
    "COVID-19 has many symptoms.",
    "COVID-19 symptoms are bad.",
    "COVID-19 symptoms are not nice",
    "COVID-19 symptoms are bad. 5G capabilities include more expansive service coverage, a higher number of available connections, and lower power consumption.",
    "COVID-19 is a disease caused by a virus. The most common symptoms are fever, chills, and sore throat, but there are a range of others.",
    "COVID-19 symptoms can include: a high temperature or shivering (chills); a new, continuous cough; a loss or change to your sense of smell or taste; and many more",
    "Dementia has the following symptom: Experiencing memory loss, poor judgment, and confusion.",
    "COVID-19 has the following symptom: Experiencing memory loss, poor judgment, and confusion.",
]
# Constructing the embed_input object
embed_input = [
    {"content": [{"type": "text", "text": doc}]} for doc in documents
]
# Embed the documents
doc_emb = co.embed(
    inputs=embed_input,
    model="embed-v4.0",
    output_dimension=1024,
    input_type="search_document",
    embedding_types=["float"],
).embeddings.float
### STEP 2: Embed the query
# Add the user query
query = "COVID-19 Symptoms"
query_input = [{"content": [{"type": "text", "text": query}]}]
# Embed the query
query_emb = co.embed(
    inputs=query_input,
    model="embed-v4.0",
    input_type="search_query",
    output_dimension=1024,
    embedding_types=["float"],
).embeddings.float
### STEP 3: Return the most similar documents
# Calculate similarity scores
scores = np.dot(query_emb, np.transpose(doc_emb))[0]
# Sort and filter documents based on scores
top_n = 5
top_doc_idxs = np.argsort(-scores)[:top_n]
# Display search results
for idx, docs_idx in enumerate(top_doc_idxs):
    print(f"Rank: {idx+1}")
    print(f"Document: {documents[docs_idx]}\n")

Here’s a sample output:

Rank: 1
Document: COVID-19 symptoms can include: a high temperature or shivering (chills); a new, continuous cough; a loss or change to your sense of smell or taste; and many more
Rank: 2
Document: COVID-19 is a disease caused by a virus. The most common symptoms are fever, chills, and sore throat, but there are a range of others.
Rank: 3
Document: COVID-19 has the following symptom: Experiencing memory loss, poor judgment, and confusion.
Rank: 4
Document: COVID-19 has many symptoms.
Rank: 5
Document: COVID-19 symptoms are not nice

The Embed endpoint also supports multilingual semantic search via embed-v4.0 and previous embed-multilingual-... models. This means you can perform semantic search on texts in different languages.

Specifically, you can do both multilingual and cross-lingual searches using one single model.

Specifically, you can do both multilingual and cross-lingual searches using one single model.

PYTHON
### STEP 1: Embed the documents
documents = [
    "Remboursement des frais de voyage : Gérez facilement vos frais de voyage en les soumettant via notre outil financier. Les approbations sont rapides et simples.",
    "Travailler de l'étranger : Il est possible de travailler à distance depuis un autre pays. Il suffit de coordonner avec votre responsable et de vous assurer d'être disponible pendant les heures de travail.",
    "Avantages pour la santé et le bien-être : Nous nous soucions de votre bien-être et proposons des adhésions à des salles de sport, des cours de yoga sur site et une assurance santé complète.",
    "Fréquence des évaluations de performance : Nous organisons des bilans informels tous les trimestres et des évaluations formelles deux fois par an.",
]
# Constructing the embed_input object
embed_input = [
    {"content": [{"type": "text", "text": doc}]} for doc in documents
]
# Embed the documents
doc_emb = co.embed(
    inputs=embed_input,
    model="embed-v4.0",
    output_dimension=1024,
    input_type="search_document",
    embedding_types=["float"],
).embeddings.float
### STEP 2: Embed the query
# Add the user query
query = "What's your remote-working policy?"
query_input = [{"content": [{"type": "text", "text": query}]}]
# Embed the query
query_emb = co.embed(
    inputs=query_input,
    model="embed-v4.0",
    input_type="search_query",
    output_dimension=1024,
    embedding_types=["float"],
).embeddings.float
### STEP 3: Return the most similar documents
# Calculate similarity scores
scores = np.dot(query_emb, np.transpose(doc_emb))[0]
# Sort and filter documents based on scores
top_n = 4
top_doc_idxs = np.argsort(-scores)[:top_n]
# Display search results
for idx, docs_idx in enumerate(top_doc_idxs):
    print(f"Rank: {idx+1}")
    print(f"Document: {documents[docs_idx]}\n")

Here’s a sample output:

Rank: 1
Document: Travailler de l'étranger : Il est possible de travailler à distance depuis un autre pays. Il suffit de coordonner avec votre responsable et de vous assurer d'être disponible pendant les heures de travail.
Rank: 2
Document: Avantages pour la santé et le bien-être : Nous nous soucions de votre bien-être et proposons des adhésions à des salles de sport, des cours de yoga sur site et une assurance santé complète.
Rank: 3
Document: Fréquence des évaluations de performance : Nous organisons des bilans informels tous les trimestres et des évaluations formelles deux fois par an.
Rank: 4
Document: Remboursement des frais de voyage : Gérez facilement vos frais de voyage en les soumettant via notre outil financier. Les approbations sont rapides et simples.

Handling PDF files, which often contain a mix of text, images, and layout information, presents a challenge for traditional embedding methods. This usually requires a multimodal generative model to pre-process the documents into a format that is suitable for the embedding model. This intermediate text representations can lose critical information; for example, the structure and precise content of tables or complex layouts might not be accurately rendered

Embed v4 solves this problem as it is designed to natively understand mixed-modality inputs. Embed v4 can directly process the PDF content, including text and images, in a single step. It generates a unified embedding that captures the semantic meaning derived from both the textual and visual elements.

Here’s an example of how to use the Embed endpoint to perform multimodal PDF search.

First, import the required libraries.

PYTHON
from pdf2image import convert_from_path
from io import BytesIO
import base64
import chromadb
import cohere

Next, turn a PDF file into a list of images, with one image per page. Then format these images into the content structure expected by the Embed endpoint.

PYTHON
pdf_path = "PDF_FILE_PATH"  # https://github.com/cohere-ai/cohere-developer-experience/raw/main/notebooks/guide/embed-v4-pdf-search/data/Samsung_Home_Theatre_HW-N950_ZA_FullManual_02_ENG_180809_2.pdf
pages = convert_from_path(pdf_path, dpi=200)
input_array = []
for page in pages:
    buffer = BytesIO()
    page.save(buffer, format="PNG")
    base64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
    base64_image = f"data:image/png;base64,{base64_str}"
    page_entry = {
        "content": [
            {"type": "text", "text": f"{pdf_path}"},
            {"type": "image_url", "image_url": {"url": base64_image}},
        ]
    }
    input_array.append(page_entry)

Next, generate the embeddings for these pages and store them in a vector database (in this example, we use Chroma).

PYTHON
# Generate the document embeddings
embeddings = []
for i in range(0, len(input_array)):
    res = co.embed(
        model="embed-v4.0",
        input_type="search_document",
        embedding_types=["float"],
        inputs=[input_array[i]],
    ).embeddings.float[0]
    embeddings.append(res)
# Store the embeddings in a vector database
ids = []
for i in range(0, len(input_array)):
    ids.append(str(i))
chroma_client = chromadb.Client()
collection = chroma_client.create_collection("pdf_pages")
collection.add(
    embeddings=embeddings,
    ids=ids,
)

Finally, provide a query and run a search over the documents. This will return a list of sorted IDs representing the most similar pages to the query.

PYTHON
query = "Do the speakers come with an optical cable?"
# Generate the query embedding
query_embeddings = co.embed(
    model="embed-v4.0",
    input_type="search_query",
    embedding_types=["float"],
    texts=[query],
).embeddings.float[0]
# Search the vector database
results = collection.query(
    query_embeddings=[query_embeddings],
    n_results=5,  # Define the top_k value
)
# Print the id of the top-ranked page
print(results["ids"][0][0])
22

The top-ranked page is shown below:

For a more complete example of multimodal PDF search, see the cookbook version.