-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathembeddings.py
41 lines (30 loc) · 1.24 KB
/
embeddings.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
from ibm_watsonx_ai.foundation_models import Embeddings
from ibm_watsonx_ai import APIClient
from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames as EmbedParams
from ibm_watsonx_ai.foundation_models.utils.enums import EmbeddingTypes
def get_embeddings(text: str):
my_credentials = {
"url": "https://us-south.ml.cloud.ibm.com",
# "apikey": 'SKILLS-NETWORK' -- leave this commented out, we'll automatically inject credentials for you
}
model_id = EmbeddingTypes.IBM_SLATE_30M_ENG
project_id = "skills-network"
space_id = None
verify = False
# Set the truncate_input_tokens to a value that is equal to or less than the maximum allowed tokens for the embedding model that you are using. If you don't specify this value and the input has more tokens than the model can process, an error is generated.
embed_params = {
EmbedParams.TRUNCATE_INPUT_TOKENS: 128,
EmbedParams.RETURN_OPTIONS: {
'input_text': True
}
}
embedding = Embeddings(
model_id=model_id,
credentials=my_credentials,
params=embed_params,
project_id=project_id,
space_id=space_id,
verify=verify
)
embedding_vectors = embedding.embed_documents(texts=[text])[0]
return embedding_vectors