4-3 Gemini API: Embeddings 快速入門指南

16th鐵人賽

Wolke

2024-08-16 06:48:32

1349 瀏覽

分享至

4-3 Gemini API: Embeddings 快速入門指南

Gemini API 生成最先進的文本嵌入。嵌入是一組浮點數，代表一個單詞、句子或段落的含義。您可以在許多下游應用中使用嵌入，例如文檔搜索。

此筆記本提供了一些快速的代碼示例，展示了如何開始生成嵌入。

安裝必要套件

!pip install -q -U google-generativeai

導入庫

import google.generativeai as genai

配置您的 API 密鑰

要運行以下代碼單元，您的 API 密鑰必須存儲在名為 GOOGLE_API_KEY 的 Colab Secret 中。如果您尚未擁有 API 密鑰，或者不確定如何創建 Colab Secret，請參見 Authentication 示例。

from google.colab import userdata
GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')
genai.configure(api_key=GOOGLE_API_KEY)

嵌入內容

使用 models/text-embedding-004 模型調用 embed_content 方法來生成文本嵌入。

text = "Hello world"
result = genai.embed_content(model="models/text-embedding-004", content=text)

# 為了使輸出易於管理，只打印嵌入的一部分
print(str(result['embedding'])[:50], '... TRIMMED]')

檢查嵌入維度

print(len(result['embedding'])) # 嵌入的維度是 768

批量嵌入內容

您可以用一次 API 調用嵌入多個提示，以提高效率。

result = genai.embed_content(
    model="models/text-embedding-004",
    content=[
      'What is the meaning of life?',
      'How much wood would a woodchuck chuck?',
      'How does the brain work?'])

for embedding in result['embedding']:
  print(str(embedding)[:50], '... TRIMMED]')

截短嵌入

text-embedding-004 模型還支持更低的嵌入維度。指定 output_dimensionality 來截短輸出。

# 未截短
result1 = genai.embed_content(
    model="models/text-embedding-004",
    content="Hello world")

# 截短
result2 = genai.embed_content(
    model="models/text-embedding-004",
    content="Hello world",
    output_dimensionality=10)

(len(result1['embedding']), len(result2['embedding']))

指定 `task_type`

讓我們來看看 embed_content 方法所接受的所有參數。共有五個參數：

model: 必填。必須是 models/text-embedding-004 或 models/embedding-001。
content: 必填。您想嵌入的內容。
task_type: 可選。嵌入將用於的任務類型。
title: 可選。如果您的任務類型是 retrieval_document（或 document），則應僅設置此參數。
output_dimensionality: 可選。輸出嵌入的縮減維度。如果設置，將從結尾截短輸出嵌入中的多餘值。這由 models/text-embedding-004 支持，但不能在 models/embedding-001 中指定。

task_type 是一個可選參數，向 API 提供您打算在應用中使用嵌入的提示。

以下是接受的 task_type 參數：

unspecified: 如果您未設置值，默認為 retrieval_query。
retrieval_query（或 query）：給定文本在搜索/檢索設置中是查詢。
retrieval_document（或 document）：給定文本是被檢索語料庫中的文檔。可選的，還設置文檔標題參數 title。
semantic_similarity（或 similarity）：給定文本將用於語義文本相似性（STS）。
classification: 給定文本將被分類。
clustering: 嵌入將用於聚類。
question_answering: 給定文本將用於問答。
fact_verification: 給定文本將用於事實驗證。

# 注意，API 根據 `task_type` 返回不同的嵌入
result1 = genai.embed_content(
    model="models/text-embedding-004",
    content="Hello world")

result2 = genai.embed_content(
    model="models/text-embedding-004",
    content="Hello world",
    task_type="document")

print(str(result1['embedding'])[:50], '... TRIMMED]')
print(str(result2['embedding'])[:50], '... TRIMMED]')