mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 23:02:28 +00:00
Jack Ye
b21fb2a4dc
The pages describing this code lived in lancedb/docs, a repository with ten stars that neither contributing guide mentioned. A user who hit a wrong page while reading the source had to find a separate, near-invisible repository to fix it. These 102 pages now sit beside the code, and a documentation fix is a pull request in the repository the reader already had open. `docs/web` is a complete Mintlify site: `cd docs/web && mint dev` serves it with no other checkout. Enterprise pages stay private and are merged in when the published site is assembled, so links into them do not resolve here — every open-source page does. `docs/web-tests` holds the tests the code examples are extracted from. They ran against a published wheel before; here they can run against this repository's own SDK, so an example that stops compiling fails the pull request that broke it rather than being noticed later. Both directories sit beside `docs/src` and are invisible to it: mkdocs reads `docs_dir: src`, so the SDK reference builds and deploys exactly as before.
142 lines
4.6 KiB
Python
142 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import lancedb
|
|
|
|
# --8<-- [start:imports]
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
from lancedb.embeddings import get_registry
|
|
|
|
# --8<-- [end:imports]
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.skip(reason="OpenAI is not available in the test environment")
|
|
def test_create_embedding_function():
|
|
# --8<-- [start:create_embedding_function]
|
|
func = get_registry().get("openai").create(
|
|
name="text-embedding-3-small",
|
|
max_retries=7,
|
|
)
|
|
# --8<-- [end:create_embedding_function]
|
|
assert func is not None
|
|
|
|
|
|
@pytest.mark.skip(reason="OpenAI is not available in the test environment")
|
|
def test_embeddings_openai():
|
|
# --8<-- [start:openai_embeddings]
|
|
db = lancedb.connect("/tmp/db")
|
|
func = get_registry().get("openai").create(name="text-embedding-ada-002")
|
|
|
|
class Words(LanceModel):
|
|
text: str = func.SourceField()
|
|
vector: Vector(func.ndims()) = func.VectorField()
|
|
|
|
table = db.create_table("words", schema=Words, mode="overwrite")
|
|
table.add([{"text": "hello world"}, {"text": "goodbye world"}])
|
|
|
|
query = "greetings"
|
|
actual = table.search(query).limit(1).to_pydantic(Words)[0]
|
|
print(actual.text)
|
|
# --8<-- [end:openai_embeddings]
|
|
|
|
|
|
@pytest.mark.skip(reason="OpenAI is not available in the test environment")
|
|
@pytest.mark.asyncio
|
|
async def test_embeddings_openai_async():
|
|
uri = "memory://"
|
|
# --8<-- [start:async_openai_embeddings]
|
|
db = await lancedb.connect_async(uri)
|
|
func = get_registry().get("openai").create(name="text-embedding-ada-002")
|
|
|
|
class Words(LanceModel):
|
|
text: str = func.SourceField()
|
|
vector: Vector(func.ndims()) = func.VectorField()
|
|
|
|
table = await db.create_table("words", schema=Words, mode="overwrite")
|
|
await table.add([{"text": "hello world"}, {"text": "goodbye world"}])
|
|
|
|
query = "greetings"
|
|
actual = await (await table.search(query)).limit(1).to_pydantic(Words)[0]
|
|
print(actual.text)
|
|
# --8<-- [end:async_openai_embeddings]
|
|
|
|
|
|
@pytest.mark.skip(reason="OpenAI is not available in the test environment")
|
|
def test_embeddings_manual_query():
|
|
# --8<-- [start:manual_query_embeddings]
|
|
db = lancedb.connect("/tmp/db")
|
|
func = get_registry().get("openai").create(name="text-embedding-ada-002")
|
|
|
|
class Words(LanceModel):
|
|
text: str = func.SourceField()
|
|
vector: Vector(func.ndims()) = func.VectorField()
|
|
|
|
table = db.create_table("words", schema=Words, mode="overwrite")
|
|
table.add([{"text": "hello world"}, {"text": "goodbye world"}])
|
|
|
|
query_vector = func.generate_embeddings(["greetings"])[0]
|
|
# --8<-- [start:manual_query_search]
|
|
# query_vector is assumed to already be generated by your embedding function
|
|
actual = table.search(query_vector).limit(1).to_pydantic(Words)[0]
|
|
print(actual.text)
|
|
# --8<-- [end:manual_query_search]
|
|
# --8<-- [end:manual_query_embeddings]
|
|
|
|
|
|
def test_custom_embedding_function():
|
|
# --8<-- [start:embedding_function]
|
|
from functools import cached_property
|
|
|
|
from lancedb.embeddings import TextEmbeddingFunction, register
|
|
|
|
class MyEmbeddingModel:
|
|
def __init__(self, model_name: str):
|
|
self.model_name = model_name
|
|
|
|
def encode(self, texts: list[str]) -> list[list[float]]:
|
|
return [[1.0, 2.0, 3.0] for _ in texts]
|
|
|
|
@register("my-embedder")
|
|
class MyTextEmbedder(TextEmbeddingFunction):
|
|
model_name: str = "my-model"
|
|
|
|
def generate_embeddings(self, texts: list[str]) -> list[list[float]]:
|
|
# Your embedding logic here
|
|
return self._model.encode(texts)
|
|
|
|
def ndims(self) -> int:
|
|
# Return the dimensionality of the embeddings
|
|
return len(self.generate_embeddings(["test"])[0])
|
|
|
|
@cached_property
|
|
def _model(self) -> MyEmbeddingModel:
|
|
# Initialize your model once
|
|
return MyEmbeddingModel(self.model_name)
|
|
# --8<-- [end:embedding_function]
|
|
|
|
|
|
def test_embeddings_secret():
|
|
# --8<-- [start:register_secret]
|
|
registry = get_registry()
|
|
registry.set_var("api_key", "sk-...")
|
|
|
|
func = registry.get("openai").create(api_key="$var:api_key")
|
|
# --8<-- [end:register_secret]
|
|
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
pytest.skip("torch not installed")
|
|
|
|
# --8<-- [start:register_device]
|
|
import torch
|
|
|
|
registry = get_registry()
|
|
if torch.cuda.is_available():
|
|
registry.set_var("device", "cuda")
|
|
|
|
func = registry.get("huggingface").create(device="$var:device:cpu")
|
|
# --8<-- [end:register_device]
|
|
assert func.device == "cuda" if torch.cuda.is_available() else "cpu"
|