Generative AI Google Brain

embeddings-task-type-embedding

3 min read

Using “task type” embeddings for improving RAG search quality

Source notebook

Repo path: embeddings/task-type-embedding.ipynb · Open on GitHub · intermediate

Shows how task-type embeddings improve RAG Q&A retrieval quality and MRR.

Summary

This notebook explains why plain semantic similarity can fail for RAG because questions and answers are not always semantically similar. It demonstrates generating embeddings with gemini-embedding-001 using task types such as SEMANTIC_SIMILARITY, QUESTION_ANSWERING, and RETRIEVAL_DOCUMENT. It then evaluates 1K sampled NQ-Open question-answer pairs by computing cosine similarities, top-100 ranks, and MRR, showing better Q&A retrieval with task-type embeddings.

Key code patterns

Initialize GenAI client for Vertex AI

from google import genai
 
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

Creates the SDK client used to call the embedding model through Vertex AI.

Embed with task type

from google.genai.types import EmbedContentConfig
 
MODEL_ID = "gemini-embedding-001"
response = client.models.embed_content(
    model=MODEL_ID,
    contents=contents,
    config=EmbedContentConfig(task_type=task_type),
)

Task type controls how the embedding space is optimized for retrieval behavior.

Q&A retrieval pairing

q_emb = get_embeddings([QUESTION], "QUESTION_ANSWERING")
a1_emb = get_embeddings([ANSWER_1], "RETRIEVAL_DOCUMENT")
a2_emb = get_embeddings([ANSWER_2], "RETRIEVAL_DOCUMENT")
print(cosine_similarity(q_emb, a1_emb))
print(cosine_similarity(q_emb, a2_emb))

Uses different task types for questions and answer documents to improve answer matching.

Batch embedding generation

for i in tqdm(range(0, len(df), 5)):
    q_embeddings.extend(get_embeddings(df.loc[i:i+4, "question"].tolist(), question_task_type))
    answers = [row[0] for row in df.loc[i:i+4, "answer"]]
    a_embeddings.extend(get_embeddings(answers, answer_task_type))

Bundles five rows per API call and stores question and answer embeddings in the dataframe.

MRR evaluation

q_and_a_similarities = cosine_similarity(q_embeddings, a_embeddings)
sorted_index = sorted(sim_index, key=lambda x: x[1], reverse=True)[:100]
reciprocal_ranks.append(1 / (i + 1))

Measures retrieval quality by ranking the ground-truth answer among the top similar answers.

Models & APIs used

  • Models: gemini-embedding-001
  • APIs / services: Vertex AI, Agent Platform Embeddings API
  • SDKs / libraries: google-genai, pandas, scikit-learn, tqdm

When to use this

Use this pattern when improving RAG answer retrieval quality with embeddings before adding LLM-based query expansion or custom tuning.

Gotchas & caveats

  • A Google Cloud project is required and aiplatform.googleapis.com must be enabled.
  • Colab requires explicit user authentication with google.colab.auth.authenticate_user().
  • LOCATION defaults to us-central1 unless GOOGLE_CLOUD_REGION is set.
  • Embedding 1K rows takes a few minutes; increasing to 10K rows may take about 30 minutes.
  • Task-type embeddings are trained on generic web corpus and may be less effective for proprietary documents or unknown topics.

Best practices

  • Use QUESTION_ANSWERING for question texts and RETRIEVAL_DOCUMENT for answer documents in Q&A RAG.
  • Use SEMANTIC_SIMILARITY as a baseline when comparing retrieval quality.
  • Evaluate search quality with ranking metrics such as Mean Reciprocal Rank.
  • Batch multiple texts per embedding API call to reduce repeated calls.
  • Consider tuning text embeddings when pre-trained task-type embeddings do not fit proprietary or specialized content.

Graph View

Backlinks