mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 12:12:27 +00:00
Will JonesandClaude Opus 5.5
9020640538
BREAKING CHANGE: `SentenceTransformerEmbeddings.trust_remote_code` now
defaults to `False` (was `True`). When an embedding function is restored
from table metadata, `trust_remote_code=True` and custom endpoints
(OpenAI `base_url`, Ollama `host`/proxy, watsonx `url`) are only honored
if the process approved them; otherwise computing embeddings raises
`UntrustedEmbeddingConfigError`. Functions created in the same process
are approved automatically. To reopen such a table in a new process,
call `get_registry().allow_remote_code(model, revision=...)` or
`get_registry().allow_endpoint(url)` first.
## Summary
Embedding functions are rebuilt from table metadata, which anyone who
can write the table controls. Since sentence-transformers defaulted
`trust_remote_code` to `True`, opening and querying a crafted table ran
arbitrary code from a model repository (GHSA-x8v2-vj73-gfwm). Metadata
could likewise point `base_url` at another server, sending it queries
and the `OPENAI_API_KEY` from the environment.
The table can no longer grant itself trust:
```python
registry = get_registry()
registry.allow_remote_code("nomic-ai/nomic-embed-text-v1.5", revision="<commit>")
registry.allow_endpoint("http://my-ollama-host:11434")
```
Trust is keyed by model name (optionally pinned to a revision), so
approving one model doesn't let a table swap in another. Restored
functions that aren't approved are replaced by a placeholder that raises
on use without constructing the real function (`huggingface` loads its
model in `__init__`), so non-embedding operations on the table keep
working. `sentence-transformers` and `huggingface` gain a `revision`
parameter. Custom embedding functions can opt in by overriding the new
`remote_code_model()` / `endpoints()` hooks.
Not covered: TypeScript (its built-ins have neither remote code nor
configurable endpoints), Bedrock `assumed_role`, and third-party
functions that don't implement the hooks.
## Evidence
`test_embeddings_trust.py` covers restore with and without approval,
revision pinning, swapped model names, endpoint gating, and that the
table stays usable. Opening a table with crafted metadata and searching
it now gives:
```
UntrustedEmbeddingConfigError: Embedding function 'sentence-transformers' was restored from table metadata with settings this process has not approved:
- runs code from model 'attacker/model' at revision None. If you trust it, call get_registry().allow_remote_code('attacker/model', revision=None)
```
## Merge Risk
**Door:** One-way: changes a default and makes previously working
reopen-and-search flows require explicit approval.
**Blast Radius:** public-api (Python embeddings); users of remote-code
models or custom endpoints who reopen tables in a new process.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
|
|
from functools import cached_property
|
|
from typing import List, Any, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from pydantic import ConfigDict, PrivateAttr
|
|
|
|
from ..util import attempt_import_or_raise
|
|
from .base import EmbeddingFunction
|
|
from .registry import register
|
|
from .utils import TEXT
|
|
|
|
|
|
@register("huggingface")
|
|
class TransformersEmbeddingFunction(EmbeddingFunction):
|
|
"""
|
|
An embedding function that can use any model from the transformers library.
|
|
|
|
Parameters:
|
|
----------
|
|
name : str
|
|
The name of the model to use. This should be a model name that can be loaded
|
|
by transformers.AutoModel.from_pretrained. For example, "bert-base-uncased".
|
|
default: "colbert-ir/colbertv2.0""
|
|
device : str
|
|
The device to use for the model. Default is "cpu".
|
|
show_progress_bar : bool
|
|
Whether to show a progress bar when loading the model. Default is True.
|
|
trust_remote_code : bool
|
|
Whether or not to allow for custom models defined on the HuggingFace
|
|
Hub in their own modeling files. This option should only be set to True
|
|
for repositories you trust and in which you have read the code, as it
|
|
will execute code present on the Hub on your local machine. When
|
|
restored from table metadata, this is only enabled for models allowed with
|
|
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
|
|
revision : Optional[str]
|
|
The model revision to load, such as a commit hash. Default is None.
|
|
|
|
to download package, run :
|
|
`pip install transformers`
|
|
you may need to install pytorch as well - `https://pytorch.org/get-started/locally/`
|
|
|
|
"""
|
|
|
|
name: str = "colbert-ir/colbertv2.0"
|
|
device: str = "cpu"
|
|
trust_remote_code: bool = False
|
|
revision: Optional[str] = None
|
|
_tokenizer: Any = PrivateAttr()
|
|
_model: Any = PrivateAttr()
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._ndims = None
|
|
transformers = attempt_import_or_raise("transformers")
|
|
self._tokenizer = transformers.AutoTokenizer.from_pretrained(
|
|
self.name, revision=self.revision
|
|
)
|
|
self._model = transformers.AutoModel.from_pretrained(
|
|
self.name, trust_remote_code=self.trust_remote_code, revision=self.revision
|
|
)
|
|
self._model.to(self.device)
|
|
|
|
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
|
|
return (self.name, self.revision)
|
|
|
|
model_config = ConfigDict(ignored_types=(cached_property,))
|
|
|
|
def ndims(self):
|
|
self._ndims = self._model.config.hidden_size
|
|
return self._ndims
|
|
|
|
def compute_query_embeddings(self, query: str, *args, **kwargs) -> List[np.array]:
|
|
return self.compute_source_embeddings(query)
|
|
|
|
def compute_source_embeddings(self, texts: TEXT, *args, **kwargs) -> List[np.array]:
|
|
texts = self.sanitize_input(texts)
|
|
embedding = []
|
|
for text in texts:
|
|
encoding = self._tokenizer(
|
|
text, return_tensors="pt", padding=True, truncation=True
|
|
).to(self.device)
|
|
emb = self._model(**encoding).last_hidden_state.mean(dim=1).squeeze()
|
|
embedding.append(emb.tolist())
|
|
|
|
return embedding
|
|
|
|
|
|
@register("colbert")
|
|
class ColbertEmbeddings(TransformersEmbeddingFunction):
|
|
"""
|
|
An embedding function that uses the colbert model from the huggingface library.
|
|
|
|
Parameters:
|
|
----------
|
|
name : str
|
|
The name of the model to use. This should be a model name that can be loaded
|
|
by transformers.AutoModel.from_pretrained. For example, "bert-base-uncased".
|
|
default: "colbert-ir/colbertv2.0""
|
|
|
|
to download package, run :
|
|
`pip install transformers`
|
|
you may need to install pytorch as well - `https://pytorch.org/get-started/locally/`
|
|
|
|
"""
|
|
|
|
name: str = "colbert-ir/colbertv2.0"
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|