llama-index-embeddings-ibm

LlamaIndex Embeddings Integration: IBM

This package integrates the LlamaIndex LLMs API with the IBM watsonx.ai Foundation Models API by leveraging ibm-watsonx-ai SDK. With this integration, you can use one of the embedding models that are available in IBM watsonx.ai to embed a single string or a list of strings.

Installation

pip install llama-index-embeddings-ibm

Usage

Setting up

To use IBM's models, you must have an IBM Cloud user API key. Here's how to obtain and set up your API key:

• Obtain an API Key: For more details on how to create and manage an API key, refer to Managing user API keys.
• Set the API Key as an Environment Variable: For security reasons, it's recommended to not hard-code your API key directly in your scripts. Instead, set it up as an environment variable. You can use the following code to prompt for the API key and set it as an environment variable:

import os
from getpass import getpass

watsonx_api_key = getpass()
os.environ["WATSONX_APIKEY"] = watsonx_api_key

Alternatively, you can set the environment variable in your terminal.

• Linux/macOS: Open your terminal and execute the following command:

  export WATSONX_APIKEY='your_ibm_api_key'
  

  To make this environment variable persistent across terminal sessions, add the above line to your ~/.bashrc, ~/.bash_profile, or ~/.zshrc file.

• Windows: For Command Prompt, use:

  set WATSONX_APIKEY=your_ibm_api_key
  

Load the model

You might need to adjust embedding parameters for different tasks.

truncate_input_tokens = 3

Initialize the WatsonxEmbeddings class with the previously set parameters.

Note:

• To provide context for the API call, you must pass the project_id or space_id. To get your project or space ID, open your project or space, go to the Manage tab, and click General. For more information see: Project documentation or Deployment space documentation.
• Depending on the region of your provisioned service instance, use one of the URLs listed in watsonx.ai API Authentication.

In this example, we’ll use the project_id and Dallas URL.

You need to specify the model_id that will be used for inferencing.

from llama_index.embeddings.ibm import WatsonxEmbeddings

watsonx_embedding = WatsonxEmbeddings(
    model_id="ibm/slate-125m-english-rtrvr",
    url="https://us-south.ml.cloud.ibm.com",
    project_id="PASTE YOUR PROJECT_ID HERE",
    truncate_input_tokens=truncate_input_tokens,
)

Alternatively, you can use Cloud Pak for Data credentials. For details, see watsonx.ai software setup.

watsonx_embedding = WatsonxEmbeddings(
    model_id="ibm/slate-125m-english-rtrvr",
    url="PASTE YOUR URL HERE",
    username="PASTE YOUR USERNAME HERE",
    password="PASTE YOUR PASSWORD HERE",
    instance_id="openshift",
    version="4.8",
    project_id="PASTE YOUR PROJECT_ID HERE",
    truncate_input_tokens=truncate_input_tokens,
)

Usage

Embed query

query = "Example query."

query_result = watsonx_embedding.get_query_embedding(query)
print(query_result[:5])

Embed list of texts

texts = ["This is a content of one document", "This is another document"]

doc_result = watsonx_embedding.get_text_embedding_batch(texts)
print(doc_result[0][:5])