@@ -15,18 +15,64 @@ def _split_texts_into_batches(texts: List[str], batch_size: int) -> Iterator[Lis
15
15
16
16
17
17
class AI21Embeddings (Embeddings , AI21Base ):
18
- """AI21 embedding model.
18
+ """AI21 embedding model integration .
19
19
20
- To use, you should have the 'AI21_API_KEY' environment variable set
21
- or pass as a named parameter to the constructor.
20
+ Install ``langchain_ai21`` and set environment variable ``AI21_API_KEY``.
22
21
23
- Example:
22
+ .. code-block:: bash
23
+
24
+ pip install -U langchain_ai21
25
+ export AI21_API_KEY="your-api-key"
26
+
27
+ Key init args — client params:
28
+ api_key: Optional[SecretStr]
29
+ batch_size: int
30
+ The number of texts that will be sent to the API in each batch.
31
+ Use larger batch sizes if working with many short texts. This will reduce
32
+ the number of API calls made, and can improve the time it takes to embed
33
+ a large number of texts.
34
+ num_retries: Optional[int]
35
+ Maximum number of retries for API requests before giving up.
36
+ timeout_sec: Optional[float]
37
+ Timeout in seconds for API requests. If not set, it will default to the
38
+ value of the environment variable `AI21_TIMEOUT_SEC` or 300 seconds.
39
+
40
+ See full list of supported init args and their descriptions in the params section.
41
+
42
+ Instantiate:
24
43
.. code-block:: python
25
44
26
45
from langchain_ai21 import AI21Embeddings
27
46
28
- embeddings = AI21Embeddings()
29
- query_result = embeddings.embed_query("Hello embeddings world!")
47
+ embed = AI21Embeddings(
48
+ # api_key="...",
49
+ # batch_size=128,
50
+ )
51
+
52
+ Embed single text:
53
+ .. code-block:: python
54
+
55
+ input_text = "The meaning of life is 42"
56
+ vector = embed.embed_query(input_text)
57
+ print(vector[:3])
58
+
59
+ .. code-block:: python
60
+
61
+ [-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
62
+
63
+ Embed multiple texts:
64
+ .. code-block:: python
65
+
66
+ input_texts = ["Document 1...", "Document 2..."]
67
+ vectors = embed.embed_documents(input_texts)
68
+ print(len(vectors))
69
+ # The first 3 coordinates for the first vector
70
+ print(vectors[0][:3])
71
+
72
+ .. code-block:: python
73
+
74
+ 2
75
+ [-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
30
76
"""
31
77
32
78
batch_size : int = _DEFAULT_BATCH_SIZE
0 commit comments