mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-28 00:48:40 +00:00
Will Jones
7ac5f74c80
BREAKING CHANGE: embedding function implementations in Node need to now call `resolveVariables()` in their constructors and should **not** implement `toJSON()`. This tries to address the handling of secrets. In Node, they are currently lost. In Python, they are currently leaked into the table schema metadata. This PR introduces an in-memory variable store on the function registry. It also allows embedding function definitions to label certain config values as "sensitive", and the preprocessing logic will raise an error if users try to pass in hard-coded values. Closes #2110 Closes #521 --------- Co-authored-by: Weston Pace <weston.pace@gmail.com>
108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
|
|
import os
|
|
from functools import cached_property
|
|
from typing import List, Optional, Dict, Union
|
|
|
|
from ..util import attempt_import_or_raise
|
|
from .base import TextEmbeddingFunction
|
|
from .registry import register
|
|
|
|
import numpy as np
|
|
|
|
DEFAULT_WATSONX_URL = "https://us-south.ml.cloud.ibm.com"
|
|
|
|
MODELS_DIMS = {
|
|
"ibm/slate-125m-english-rtrvr": 768,
|
|
"ibm/slate-30m-english-rtrvr": 384,
|
|
"sentence-transformers/all-minilm-l12-v2": 384,
|
|
"intfloat/multilingual-e5-large": 1024,
|
|
}
|
|
|
|
|
|
@register("watsonx")
|
|
class WatsonxEmbeddings(TextEmbeddingFunction):
|
|
"""
|
|
API Docs:
|
|
---------
|
|
https://cloud.ibm.com/apidocs/watsonx-ai#text-embeddings
|
|
|
|
Supported embedding models:
|
|
---------------------------
|
|
https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx
|
|
"""
|
|
|
|
name: str = "ibm/slate-125m-english-rtrvr"
|
|
api_key: Optional[str] = None
|
|
project_id: Optional[str] = None
|
|
url: Optional[str] = None
|
|
params: Optional[Dict] = None
|
|
|
|
@staticmethod
|
|
def sensitive_keys():
|
|
return ["api_key"]
|
|
|
|
@staticmethod
|
|
def model_names():
|
|
return [
|
|
"ibm/slate-125m-english-rtrvr",
|
|
"ibm/slate-30m-english-rtrvr",
|
|
"sentence-transformers/all-minilm-l12-v2",
|
|
"intfloat/multilingual-e5-large",
|
|
]
|
|
|
|
def ndims(self):
|
|
return self._ndims
|
|
|
|
@cached_property
|
|
def _ndims(self):
|
|
if self.name not in MODELS_DIMS:
|
|
raise ValueError(f"Unknown model name {self.name}")
|
|
return MODELS_DIMS[self.name]
|
|
|
|
def generate_embeddings(
|
|
self,
|
|
texts: Union[List[str], np.ndarray],
|
|
*args,
|
|
**kwargs,
|
|
) -> List[List[float]]:
|
|
return self._watsonx_client.embed_documents(
|
|
texts=list(texts),
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
|
|
@cached_property
|
|
def _watsonx_client(self):
|
|
ibm_watsonx_ai = attempt_import_or_raise("ibm_watsonx_ai")
|
|
ibm_watsonx_ai_foundation_models = attempt_import_or_raise(
|
|
"ibm_watsonx_ai.foundation_models"
|
|
)
|
|
|
|
kwargs = {"model_id": self.name}
|
|
if self.params:
|
|
kwargs["params"] = self.params
|
|
if self.project_id:
|
|
kwargs["project_id"] = self.project_id
|
|
elif "WATSONX_PROJECT_ID" in os.environ:
|
|
kwargs["project_id"] = os.environ["WATSONX_PROJECT_ID"]
|
|
else:
|
|
raise ValueError("WATSONX_PROJECT_ID must be set or passed")
|
|
|
|
creds_kwargs = {}
|
|
if self.api_key:
|
|
creds_kwargs["api_key"] = self.api_key
|
|
elif "WATSONX_API_KEY" in os.environ:
|
|
creds_kwargs["api_key"] = os.environ["WATSONX_API_KEY"]
|
|
else:
|
|
raise ValueError("WATSONX_API_KEY must be set or passed")
|
|
if self.url:
|
|
creds_kwargs["url"] = self.url
|
|
else:
|
|
creds_kwargs["url"] = DEFAULT_WATSONX_URL
|
|
kwargs["credentials"] = ibm_watsonx_ai.Credentials(**creds_kwargs)
|
|
|
|
return ibm_watsonx_ai_foundation_models.Embeddings(**kwargs)
|