Google GenAI 임베딩

Google GenAI 임베딩 (Google GenAI Embeddings)

Google의 google-genai 패키지를 사용해, LlamaIndex가 GoogleGenAIEmbedding 클래스를 제공해요. 최신 gemini-embedding-2-preview 모델로 Gemini와 Vertex AI 두 API 모두에서 텍스트를 임베딩할 수 있어요.

출처: 문서

본문

colab에서 이 노트북을 열고 있다면 LlamaIndex 🦙를 설치해야 할 거예요.

%pip install llama-index-embeddings-google-genai
import os


os.environ["GOOGLE_API_KEY"] = "..."

설정 (Setup)

GoogleGenAIEmbedding은 google-genai 패키지를 감싼 래퍼로, 기본적으로 Gemini와 Vertex AI 두 API를 모두 지원해요.

api_key를 직접 넘기거나, vertexai_config를 넘겨 Vertex AI API를 사용할 수 있어요.

그 외 옵션으로 embed_batch_size, model_name, embedding_config가 있어요.

from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
from google.genai.types import EmbedContentConfig


embed_model = GoogleGenAIEmbedding(
    model_name="gemini-embedding-2-preview",
    embed_batch_size=100,
    # api key를 직접 넘길 수도 있고
    # api_key="...",
    # vertexai_config를 넘길 수도 있어요
    # vertexai_config={
    #     "project": "...",
    #     "location": "...",
    # }
    # embedding_config도 넘길 수 있어요
    # embedding_config=EmbedContentConfig(...)
)

사용법 (Usage)

동기 (Sync)

embeddings = embed_model.get_text_embedding("Google Gemini Embeddings.")
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
[0.031099992, 0.02192731, -0.06523498, 0.016788177, 0.0392835]
Dimension of embeddings: 768
embeddings = embed_model.get_query_embedding("Query Google Gemini Embeddings.")
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
[0.022199392, 0.03671178, -0.06874573, 0.02195774, 0.05475164]
Dimension of embeddings: 768
embeddings = embed_model.get_text_embedding_batch(
    [
        "Google Gemini Embeddings.",
        "Google is awesome.",
        "Llamaindex is awesome.",
    ]
)
print(f"Got {len(embeddings)} embeddings")
print(f"Dimension of embeddings: {len(embeddings[0])}")
Got 3 embeddings
Dimension of embeddings: 768

비동기 (Async)

embeddings = await embed_model.aget_text_embedding("Google Gemini Embeddings.")
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
[0.031099992, 0.02192731, -0.06523498, 0.016788177, 0.0392835]
Dimension of embeddings: 768
embeddings = await embed_model.aget_query_embedding(
    "Query Google Gemini Embeddings."
)
print(embeddings[:5])
print(f"Dimension of embeddings: {len(embeddings)}")
[0.022199392, 0.03671178, -0.06874573, 0.02195774, 0.05475164]
Dimension of embeddings: 768
embeddings = await embed_model.aget_text_embedding_batch(
    [
        "Google Gemini Embeddings.",
        "Google is awesome.",
        "Llamaindex is awesome.",
    ]
)
print(f"Got {len(embeddings)} embeddings")
print(f"Dimension of embeddings: {len(embeddings[0])}")
Got 3 embeddings
Dimension of embeddings: 768