Back to Llama Index

Interacting with Embeddings deployed in Vertex AI Endpoint with LlamaIndex

docs/examples/embeddings/vertex_embedding_endpoint.ipynb

0.14.212.9 KB
Original Source

<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/embeddings/sagemaker_embedding_endpoint.ipynb" target="_parent"></a>

Interacting with Embeddings deployed in Vertex AI Endpoint with LlamaIndex

A Vertex AI endpoint is a managed resource that enables the deployment of machine learning models, such as embeddings, for making predictions on new data.

This notebook demonstrates how to interact with embedding endpoints using the VertexEndpointEmbedding class, leveraging the LlamaIndex.

Setting Up

If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

python
%pip install llama-index-embeddings-vertex-endpoint
python
! pip install llama-index

You need to specify the endpoint information (endpoint ID, project ID, and region) to interact with the model deployed in Vertex AI.

python
ENDPOINT_ID = "<-YOUR-ENDPOINT-ID->"
PROJECT_ID = "<-YOUR-PROJECT-ID->"
LOCATION = "<-YOUR-GCP-REGION->"

Credentials should be provided to connect to the endpoint. You can either:

  • Use a service account JSON file by specifying the service_account_file parameter.
  • Provide the service account information directly through the service_account_info parameter.

Example using a service account file:

python
from llama_index.embeddings.vertex_endpoint import VertexEndpointEmbedding

SERVICE_ACCOUNT_FILE = "<-YOUR-SERVICE-ACCOUNT-FILE-PATH->.json"

embed_model = VertexEndpointEmbedding(
    endpoint_id=ENDPOINT_ID,
    project_id=PROJECT_ID,
    location=LOCATION,
    service_account_file=SERVICE_ACCOUNT_FILE,
)

Example using direct service account info::

python
from llama_index.embeddings.vertex_endpoint import VertexEndpointEmbedding

SERVICE_ACCOUNT_INFO = {
    "private_key": "<-PRIVATE-KEY->",
    "client_email": "<-SERVICE-ACCOUNT-EMAIL->",
    "token_uri": "https://oauth2.googleapis.com/token",
}

embed_model = VertexEndpointEmbedding(
    endpoint_id=ENDPOINT_ID,
    project_id=PROJECT_ID,
    location=LOCATION,
    service_account_info=SERVICE_ACCOUNT_INFO,
)

Basic Usage

Call get_text_embedding

python
embeddings = embed_model.get_text_embedding(
    "Vertex AI is a managed machine learning (ML) platform provided by Google Cloud. It allows data scientists and developers to build, deploy, and scale machine learning models efficiently, leveraging Google's ML infrastructure."
)
python
embeddings[:10]

Call get_text_embedding_batch

python
embeddings = embed_model.get_text_embedding_batch(
    [
        "Vertex AI is a managed machine learning (ML) platform provided by Google Cloud. It allows data scientists and developers to build, deploy, and scale machine learning models efficiently, leveraging Google's ML infrastructure.",
        "Vertex is integrated with llamaIndex",
    ]
)
python
len(embeddings)