24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133 | class GradientEmbedding(BaseEmbedding):
"""GradientAI embedding models.
This class provides an interface to generate embeddings using a model
deployed in Gradient AI. At the initialization it requires a model_id
of the model deployed in the cluster.
Note:
Requires `gradientai` package to be available in the PYTHONPATH. It can be installed with
`pip install gradientai`.
"""
embed_batch_size: int = Field(default=GRADIENT_EMBED_BATCH_SIZE, gt=0)
_gradient: Any = PrivateAttr()
_model: Any = PrivateAttr()
@classmethod
def class_name(cls) -> str:
return "GradientEmbedding"
def __init__(
self,
*,
embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
gradient_model_slug: str,
gradient_access_token: Optional[str] = None,
gradient_workspace_id: Optional[str] = None,
gradient_host: Optional[str] = None,
**kwargs: Any,
):
"""Initializes the GradientEmbedding class.
During the initialization the `gradientai` package is imported. Using the access token,
workspace id and the slug of the model, the model is fetched from Gradient AI and prepared to use.
Args:
embed_batch_size (int, optional): The batch size for embedding generation. Defaults to 10,
must be > 0 and <= 100.
gradient_model_slug (str): The model slug of the model in the Gradient AI account.
gradient_access_token (str, optional): The access token of the Gradient AI account, if
`None` read from the environment variable `GRADIENT_ACCESS_TOKEN`.
gradient_workspace_id (str, optional): The workspace ID of the Gradient AI account, if `None`
read from the environment variable `GRADIENT_WORKSPACE_ID`.
gradient_host (str, optional): The host of the Gradient AI API. Defaults to None, which
means the default host is used.
Raises:
ImportError: If the `gradientai` package is not available in the PYTHONPATH.
ValueError: If the model cannot be fetched from Gradient AI.
"""
if embed_batch_size <= 0:
raise ValueError(f"Embed batch size {embed_batch_size} must be > 0.")
self._gradient = gradientai.Gradient(
access_token=gradient_access_token,
workspace_id=gradient_workspace_id,
host=gradient_host,
)
try:
self._model = self._gradient.get_embeddings_model(slug=gradient_model_slug)
except gradientai.openapi.client.exceptions.UnauthorizedException as e:
logger.error(f"Error while loading model {gradient_model_slug}.")
self._gradient.close()
raise ValueError("Unable to fetch the requested embeddings model") from e
super().__init__(
embed_batch_size=embed_batch_size, model_name=gradient_model_slug, **kwargs
)
async def _aget_text_embeddings(self, texts: List[str]) -> List[Embedding]:
"""
Embed the input sequence of text asynchronously.
"""
inputs = [{"input": text} for text in texts]
result = await self._model.aembed(inputs=inputs).embeddings
return [e.embedding for e in result]
def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]:
"""
Embed the input sequence of text.
"""
inputs = [{"input": text} for text in texts]
result = self._model.embed(inputs=inputs).embeddings
return [e.embedding for e in result]
def _get_text_embedding(self, text: str) -> Embedding:
"""Alias for _get_text_embeddings() with single text input."""
return self._get_text_embeddings([text])[0]
async def _aget_text_embedding(self, text: str) -> Embedding:
"""Alias for _aget_text_embeddings() with single text input."""
embedding = await self._aget_text_embeddings([text])
return embedding[0]
async def _aget_query_embedding(self, query: str) -> Embedding:
embedding = await self._aget_text_embeddings(
[f"{QUERY_INSTRUCTION_FOR_RETRIEVAL} {query}"]
)
return embedding[0]
def _get_query_embedding(self, query: str) -> Embedding:
return self._get_text_embeddings(
[f"{QUERY_INSTRUCTION_FOR_RETRIEVAL} {query}"]
)[0]
|