AI21Embeddings#

class langchain_ai21.embeddings.AI21Embeddings[source]#

Bases: Embeddings, AI21Base

AI21 embedding model integration.

Install langchain_ai21 and set environment variable AI21_API_KEY.

pip install -U langchain_ai21
export AI21_API_KEY="your-api-key"
Key init args — client params:

api_key: Optional[SecretStr] batch_size: int

The number of texts that will be sent to the API in each batch. Use larger batch sizes if working with many short texts. This will reduce the number of API calls made, and can improve the time it takes to embed a large number of texts.

num_retries: Optional[int]

Maximum number of retries for API requests before giving up.

timeout_sec: Optional[float]

Timeout in seconds for API requests. If not set, it will default to the value of the environment variable AI21_TIMEOUT_SEC or 300 seconds.

See full list of supported init args and their descriptions in the params section.

Instantiate:
from langchain_ai21 import AI21Embeddings

embed = AI21Embeddings(
    # api_key="...",
    # batch_size=128,
)
Embed single text:
input_text = "The meaning of life is 42"
vector = embed.embed_query(input_text)
print(vector[:3])

[-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
Embed multiple texts:
 input_texts = ["Document 1...", "Document 2..."]
vectors = embed.embed_documents(input_texts)
print(len(vectors))
# The first 3 coordinates for the first vector
print(vectors[0][:3])

2
[-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]

Create a new model by parsing and validating input data from keyword arguments.

Raises ValidationError if the input data cannot be parsed to form a valid model.

param api_host: str | None = None#

Host URL

param api_key: SecretStr | None = None#

API key for AI21 API.

Constraints:

  • type = string

  • writeOnly = True

  • format = password

param batch_size: int = 128#

Maximum number of texts to embed in each batch

param num_retries: int | None = None#

Maximum number of retries for API requests before giving up.

param timeout_sec: float | None = None#

Timeout in seconds.

If not set, it will default to the value of the environment variable AI21_TIMEOUT_SEC or 300 seconds.

async aembed_documents(texts: List[str]) → List[List[float]]#

Asynchronous Embed search docs.

Parameters:

texts (List[str]) – List of text to embed.

Returns:

List of embeddings.

Return type:

List[List[float]]

async aembed_query(text: str) → List[float]#

Asynchronous Embed query text.

Parameters:

text (str) – Text to embed.

Returns:

Embedding.

Return type:

List[float]

embed_documents(texts: List[str], *, batch_size: int | None = None, **kwargs: Any) → List[List[float]][source]#

Embed search docs.

Parameters:

  • texts (List[str]) –

  • batch_size (int | None) –

  • kwargs (Any) –

Return type:

List[List[float]]

embed_query(text: str, *, batch_size: int | None = None, **kwargs: Any) → List[float][source]#

Embed query text.

Parameters:

  • text (str) –

  • batch_size (int | None) –

  • kwargs (Any) –

Return type:

List[float]

On this page