lancedb/docs/web-tests/py/test_embedding.py

# 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]
    import lancedb

    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]
    import lancedb

    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]
    import lancedb

    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"