Notebooks
G
Google Gemini Talk To Documents With Embeddings

Talk To Documents With Embeddings

gemini-cookbookgemini-apiexamplesgemini
alph-notebooks/gemini-cookbook / Talk_to_documents_with_embeddings.ipynb
Copyright 2025 Google LLC.
[1]

Document search with embeddings

⚠️

This notebook requires paid tier rate limits to run properly.
(cf. pricing for more details).

Overview

This example demonstrates how to use the Gemini API to create embeddings so that you can perform document search. You will use the Python client library to build a word embedding that allows you to compare search strings, or questions, to document contents.

In this tutorial, you'll use embeddings to perform document search over a set of documents to ask questions related to the Google Car.

Setup

[2]

To run the following cell, your API key must be stored it in a Colab Secret named GOOGLE_API_KEY. If you don't already have an API key, or you're not sure how to create a Colab Secret, see the Authentication quickstart for an example.

[ ]

Embedding generation

In this section, you will see how to generate embeddings for a piece of text using the embeddings from the Gemini API.

See the Embeddings quickstart to learn more about the task_type parameter used below.

[2]
EMBEDDING_MODEL_ID 
embeddings=[ContentEmbedding(
  values=[
    -0.019380787,
    0.015025399,
    0.006310311,
    -0.057478663,
    0.011998727,
    <... 3067 more items ...>,
  ]
)] metadata=None

Building an embeddings database

Here are three sample texts to use to build the embeddings database. You will use the Gemini API to create embeddings of each of the documents. Turn them into a dataframe for better visualization.

[3]

Organize the contents of the dictionary into a dataframe for better visualization.

[4]

Get the embeddings for each of these bodies of text. Add this information to the dataframe.

[5]

Document search with Q&A

Now that the embeddings are generated, let's create a Q&A system to search these documents. You will ask a question about hyperparameter tuning, create an embedding of the question, and compare it against the collection of embeddings in the dataframe.

The embedding of the question will be a vector (list of float values), which will be compared against the vector of the documents using the dot product. This vector returned from the API is already normalized. The dot product represents the similarity in direction between two vectors.

The values of the dot product can range between -1 and 1, inclusive. If the dot product between two vectors is 1, then the vectors are in the same direction. If the dot product value is 0, then these vectors are orthogonal, or unrelated, to each other. Lastly, if the dot product is -1, then the vectors point in the opposite direction and are not similar to each other.

Note, with the new embeddings model (gemini-embedding-001), specify the task type as QUERY for user query and DOCUMENT when embedding a document text.

Task TypeDescription
RETRIEVAL_QUERYSpecifies the given text is a query in a search/retrieval setting.
RETRIEVAL_DOCUMENTSpecifies the given text is a document in a search/retrieval setting.
[6]

Use the find_best_passage function to calculate the dot products, and then sort the dataframe from the largest to smallest dot product value to retrieve the relevant passage out of the database.

[7]

View the most relevant document from the database:

[9]

Question and Answering Application

Let's try to use the text generation API to create a Q & A system. Input your own custom data below to create a simple question and answering example. You will still use the dot product as a metric of similarity.

[10]
[11]

Choose one of the Gemini content generation models in order to find the answer to your query.

[13]
MODEL_ID
[14]

Next steps

Check out the embeddings quickstart to learn more, and browse the cookbook for more examples.