From cfdd91b19d6dd51895cd0ec0b866e488468c33cf Mon Sep 17 00:00:00 2001 From: desmondcheongzx Date: Sun, 31 Aug 2025 13:44:20 -0700 Subject: [PATCH 01/11] impl embed_image --- daft/ai/_expressions.py | 15 +++- daft/ai/openai/__init__.py | 5 +- daft/ai/protocols.py | 18 +++++ daft/ai/provider.py | 17 ++++- daft/ai/sentence_transformers/__init__.py | 5 +- .../ai/sentence_transformers/text_embedder.py | 2 +- daft/ai/transformers/__init__.py | 33 +++++++++ daft/ai/transformers/image_embedder.py | 73 +++++++++++++++++++ daft/functions/ai/__init__.py | 18 +++++ 9 files changed, 181 insertions(+), 5 deletions(-) create mode 100644 daft/ai/transformers/__init__.py create mode 100644 daft/ai/transformers/image_embedder.py diff --git a/daft/ai/_expressions.py b/daft/ai/_expressions.py index 8cc0b133a9..558f5cb0a4 100644 --- a/daft/ai/_expressions.py +++ b/daft/ai/_expressions.py @@ -4,7 +4,7 @@ if TYPE_CHECKING: from daft import Series - from daft.ai.protocols import TextEmbedder, TextEmbedderDescriptor + from daft.ai.protocols import ImageEmbedder, ImageEmbedderDescriptor, TextEmbedder, TextEmbedderDescriptor from daft.ai.typing import Embedding @@ -19,3 +19,16 @@ def __init__(self, text_embedder: TextEmbedderDescriptor): def __call__(self, text_series: Series) -> list[Embedding]: text = text_series.to_pylist() return self.text_embedder.embed_text(text) if text else [] + + +class _ImageEmbedderExpression: + """Function expression implementation for a ImageEmbedder protocol.""" + + image_embedder: ImageEmbedder + + def __init__(self, image_embedder: ImageEmbedderDescriptor): + self.image_embedder = image_embedder.instantiate() + + def __call__(self, image_series: Series) -> list[Embedding]: + image = image_series.to_pylist() + return self.image_embedder.embed_image(image) if image else [] diff --git a/daft/ai/openai/__init__.py b/daft/ai/openai/__init__.py index 7f400cd4f1..fcf423fc97 100644 --- a/daft/ai/openai/__init__.py +++ b/daft/ai/openai/__init__.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: from daft.ai.openai.typing import OpenAIProviderOptions - from daft.ai.protocols import TextEmbedder, TextEmbedderDescriptor + from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor from daft.ai.typing import Options __all__ = [ @@ -36,3 +36,6 @@ def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmb model_name=(model or "text-embedding-3-small"), model_options=options, ) + + def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor: + raise NotImplementedError("embed_image is not currently implemented for the OpenAI provider") diff --git a/daft/ai/protocols.py b/daft/ai/protocols.py index 71cdbaca7c..af48761e98 100644 --- a/daft/ai/protocols.py +++ b/daft/ai/protocols.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from daft.ai.typing import Embedding, EmbeddingDimensions + from daft.dependencies import np @runtime_checkable @@ -24,3 +25,20 @@ class TextEmbedderDescriptor(Descriptor[TextEmbedder]): @abstractmethod def get_dimensions(self) -> EmbeddingDimensions: """Returns the dimensions of the embeddings produced by the described TextEmbedder.""" + + +@runtime_checkable +class ImageEmbedder(Protocol): + """Protocol for image embedding implementations.""" + + def embed_image(self, image: list[np.ndarray]) -> list[Embedding]: + """Embeds a batch of images into an embedding vector.""" + ... + + +class ImageEmbedderDescriptor(Descriptor[ImageEmbedder]): + """Descriptor for a ImageEmbedder implementation.""" + + @abstractmethod + def get_dimensions(self) -> EmbeddingDimensions: + """Returns the dimensions of the embeddings produced by the described ImageEmbedder.""" diff --git a/daft/ai/provider.py b/daft/ai/provider.py index 35ea43328d..41a1eb045b 100644 --- a/daft/ai/provider.py +++ b/daft/ai/provider.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from daft.ai.openai.typing import OpenAIProviderOptions - from daft.ai.protocols import TextEmbedderDescriptor + from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor class ProviderImportError(ImportError): @@ -34,9 +34,19 @@ def load_sentence_transformers(name: str | None = None, **options: Any) -> Provi raise ProviderImportError(["sentence_transformers", "torch"]) from e +def load_transformers(name: str | None = None, **options: Any) -> Provider: + try: + from daft.ai.transformers import TransformersProvider + + return TransformersProvider(name, **options) + except ImportError as e: + raise ProviderImportError(["transformers", "torch", "torchvision"]) from e + + PROVIDERS: dict[str, Callable[..., Provider]] = { "openai": load_openai, "sentence_transformers": load_sentence_transformers, + "transformers": load_transformers, } @@ -65,3 +75,8 @@ def name(self) -> str: def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor: """Returns a TextEmbedderDescriptor for this provider.""" ... + + @abstractmethod + def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor: + """Returns a ImageEmbedderDescriptor for this provider.""" + ... diff --git a/daft/ai/sentence_transformers/__init__.py b/daft/ai/sentence_transformers/__init__.py index fd675eab2c..9e8e6db3bf 100644 --- a/daft/ai/sentence_transformers/__init__.py +++ b/daft/ai/sentence_transformers/__init__.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from daft.ai.protocols import TextEmbedderDescriptor + from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor from daft.ai.typing import Options __all__ = [ @@ -28,3 +28,6 @@ def name(self) -> str: def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor: return SentenceTransformersTextEmbedderDescriptor(model or "sentence-transformers/all-MiniLM-L6-v2", options) + + def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor: + raise NotImplementedError("embed_image is not currently implemented for the Sentence Transformers provider") diff --git a/daft/ai/sentence_transformers/text_embedder.py b/daft/ai/sentence_transformers/text_embedder.py index 328cfccfff..646c3517f5 100644 --- a/daft/ai/sentence_transformers/text_embedder.py +++ b/daft/ai/sentence_transformers/text_embedder.py @@ -30,7 +30,7 @@ def get_options(self) -> Options: return self.options def get_dimensions(self) -> EmbeddingDimensions: - dimensions = AutoConfig.from_pretrained(self.model).hidden_size + dimensions = AutoConfig.from_pretrained(self.model, trust_remote_code=True).hidden_size return EmbeddingDimensions(size=dimensions, dtype=DataType.float32()) def instantiate(self) -> TextEmbedder: diff --git a/daft/ai/transformers/__init__.py b/daft/ai/transformers/__init__.py new file mode 100644 index 0000000000..6bd7e2c24e --- /dev/null +++ b/daft/ai/transformers/__init__.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from daft.ai.provider import Provider + +from daft.ai.transformers.image_embedder import TransformersImageEmbedderDescriptor +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from daft.ai.protocols import ImageEmbedderDescriptor, TextEmbedderDescriptor + from daft.ai.typing import Options + +__all__ = [ + "TransformersProvider", +] + + +class TransformersProvider(Provider): + _name: str + _options: Options + + def __init__(self, name: str | None = None, **options: Any): + self._name = name if name else "transformers" + self._options = options + + @property + def name(self) -> str: + return self._name + + def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor: + return TransformersImageEmbedderDescriptor(model or "openai/clip-vit-base-patch32", options) + + def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor: + raise NotImplementedError("embed_text is not currently implemented for the Transformers provider") diff --git a/daft/ai/transformers/image_embedder.py b/daft/ai/transformers/image_embedder.py new file mode 100644 index 0000000000..b7db8dc442 --- /dev/null +++ b/daft/ai/transformers/image_embedder.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +from transformers import AutoModel, AutoProcessor + +from daft import DataType +from daft.ai.protocols import ImageEmbedder, ImageEmbedderDescriptor +from daft.ai.typing import EmbeddingDimensions, Options +from daft.dependencies import pil_image + +if TYPE_CHECKING: + from daft.ai.typing import Embedding + from daft.dependencies import np + + +@dataclass +class TransformersImageEmbedderDescriptor(ImageEmbedderDescriptor): + model: str + options: Options + + def get_provider(self) -> str: + return "transformers" + + def get_model(self) -> str: + return self.model + + def get_options(self) -> Options: + return self.options + + def get_dimensions(self) -> EmbeddingDimensions: + from transformers import AutoConfig + + config = AutoConfig.from_pretrained(self.model, trust_remote_code=True) + # For CLIP models, the image embedding dimension is typically in projection_dim or hidden_size. + embedding_size = getattr(config, "projection_dim", getattr(config, "hidden_size", 512)) + return EmbeddingDimensions(size=embedding_size, dtype=DataType.float32()) + + def instantiate(self) -> ImageEmbedder: + return TransformersImageEmbedder(self.model, **self.options) + + +class TransformersImageEmbedder(ImageEmbedder): + model: Any + options: Options + + def __init__(self, model_name_or_path: str, **options: Any): + if not pil_image.module_available(): + raise ImportError("Pillow is required for image processing but not available") + + self.device = ( + torch.device("cuda") + if torch.cuda.is_available() + else torch.device("mps") + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + else torch.device("cpu") + ) + self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True).to(self.device) + self.processor = AutoProcessor.from_pretrained(model_name_or_path, trust_remote_code=True, use_fast=True) + self.options = options + + def embed_image(self, images: list[np.ndarray]) -> list[Embedding]: + # TODO(desmond): There's potential for image decoding and processing on the GPU with greater + # performance. Methods differ a little between different models, so let's do it later. + pil_images = [pil_image.fromarray(image) for image in images] + processed = self.processor(images=pil_images, return_tensors="pt") + pixel_values = processed["pixel_values"].to(self.device) + + with torch.inference_mode(): + embeddings = self.model.get_image_features(pixel_values) + return embeddings.cpu().numpy().tolist() diff --git a/daft/functions/ai/__init__.py b/daft/functions/ai/__init__.py index d5745114be..6bc814af89 100644 --- a/daft/functions/ai/__init__.py +++ b/daft/functions/ai/__init__.py @@ -85,3 +85,21 @@ def embed_text( ) expr = expr.with_init_args(text_embedder) return expr(text) + + +def embed_image( + image: Expression, + *, + provider: str | Provider | None = None, + model: str | None = None, + **options: str, +) -> Expression: + from daft.ai._expressions import _ImageEmbedderExpression + from daft.ai.protocols import ImageEmbedder + + image_embedder = _resolve_provider(provider, "transformers").get_image_embedder(model, **options) + expr = udf(return_dtype=image_embedder.get_dimensions().as_dtype(), concurrency=1, use_process=False)( + _ImageEmbedderExpression + ) + expr = expr.with_init_args(image_embedder) + return expr(image) From d858ca08fb8a18f0a8130e23b6fea06cf93e3b0f Mon Sep 17 00:00:00 2001 From: desmondcheongzx Date: Sun, 31 Aug 2025 14:10:42 -0700 Subject: [PATCH 02/11] add tests --- daft/ai/transformers/__init__.py | 5 ++ daft/ai/transformers/image_embedder.py | 3 - daft/functions/ai/__init__.py | 14 ++++ tests/ai/test_transformers.py | 109 +++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 3 deletions(-) create mode 100644 tests/ai/test_transformers.py diff --git a/daft/ai/transformers/__init__.py b/daft/ai/transformers/__init__.py index 6bd7e2c24e..60c06229bc 100644 --- a/daft/ai/transformers/__init__.py +++ b/daft/ai/transformers/__init__.py @@ -3,6 +3,7 @@ from daft.ai.provider import Provider from daft.ai.transformers.image_embedder import TransformersImageEmbedderDescriptor +from daft.dependencies import pil_image from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -27,6 +28,10 @@ def name(self) -> str: return self._name def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor: + # Raise an error early if PIL is not available. + if not pil_image.module_available(): + raise ImportError("Pillow is required for image processing but not available") + return TransformersImageEmbedderDescriptor(model or "openai/clip-vit-base-patch32", options) def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmbedderDescriptor: diff --git a/daft/ai/transformers/image_embedder.py b/daft/ai/transformers/image_embedder.py index b7db8dc442..34dfea46da 100644 --- a/daft/ai/transformers/image_embedder.py +++ b/daft/ai/transformers/image_embedder.py @@ -47,9 +47,6 @@ class TransformersImageEmbedder(ImageEmbedder): options: Options def __init__(self, model_name_or_path: str, **options: Any): - if not pil_image.module_available(): - raise ImportError("Pillow is required for image processing but not available") - self.device = ( torch.device("cuda") if torch.cuda.is_available() diff --git a/daft/functions/ai/__init__.py b/daft/functions/ai/__init__.py index 6bc814af89..ba9838e83c 100644 --- a/daft/functions/ai/__init__.py +++ b/daft/functions/ai/__init__.py @@ -94,6 +94,20 @@ def embed_image( model: str | None = None, **options: str, ) -> Expression: + """Returns an expression that embeds images using the specified image model and provider. + + Args: + image (Expression): The input image column expression. + provider (str | Provider | None): The provider to use for the image model. If None, the default provider is used. + model (str | None): The image model to use. Can be a model instance or a model name. If None, the default model is used. + **options: Any additional options to pass for the model. + + Note: + Make sure the required provider packages are installed (e.g. vllm, transformers, openai). + + Returns: + Expression: An expression representing the embedded image vectors. + """ from daft.ai._expressions import _ImageEmbedderExpression from daft.ai.protocols import ImageEmbedder diff --git a/tests/ai/test_transformers.py b/tests/ai/test_transformers.py new file mode 100644 index 0000000000..169e76db2e --- /dev/null +++ b/tests/ai/test_transformers.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import pytest + +pytest.importorskip("transformers") +pytest.importorskip("torch") +pytest.importorskip("PIL") + +import numpy as np +import torch + +import daft +from daft.ai.protocols import ImageEmbedderDescriptor +from daft.ai.transformers import TransformersProvider +from daft.ai.typing import EmbeddingDimensions +from daft.datatype import DataType +from daft.functions.ai import embed_image + + +def test_transformers_image_embedder_default(): + provider = TransformersProvider() + descriptor = provider.get_image_embedder() + assert isinstance(descriptor, ImageEmbedderDescriptor) + assert descriptor.get_provider() == "transformers" + assert descriptor.get_model() == "openai/clip-vit-base-patch32" + assert descriptor.get_dimensions() == EmbeddingDimensions(512, dtype=DataType.float32()) + + +@pytest.mark.parametrize( + "model_name, dimensions, run_model", + [ + ("openai/clip-vit-base-patch32", 512, True), + ("openai/clip-vit-large-patch14", 768, True), + ("openai/clip-vit-base-patch16", 512, True), + ("openai/clip-vit-large-patch14-336", 768, True), + ], +) +def test_transformers_image_embedder_other(model_name, dimensions, run_model): + mock_options = {"arg1": "val1", "arg2": "val2"} + + provider = TransformersProvider() + descriptor = provider.get_image_embedder(model_name, **mock_options) + assert isinstance(descriptor, ImageEmbedderDescriptor) + assert descriptor.get_provider() == "transformers" + assert descriptor.get_model() == model_name + assert descriptor.get_options() == mock_options + + if run_model: + embedder = descriptor.instantiate() + + # Test with a variety of image sizes and shapes that should be preprocessed correctly. + # TODO(desmond): This doesn't work with greyscale images. I wonder if we should require users to cast + # to RGB or if we want to handle this automagically. + test_image1 = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) + test_image2 = np.random.randint(0, 255, (500, 10, 3), dtype=np.uint8) + # test_image3 = np.random.randint(0, 255, (100, 100, 1), dtype=np.uint8) + test_images = [test_image1, test_image2] * 8 + + embeddings = embedder.embed_image(test_images) + assert len(embeddings) == 16 + assert len(embeddings[0]) == dimensions + else: + assert descriptor.get_dimensions() == EmbeddingDimensions(dimensions, dtype=DataType.float32()) + + +def test_transformers_device_selection(): + """Test that the embedder uses the correct device selection.""" + provider = TransformersProvider() + descriptor = provider.get_image_embedder("openai/clip-vit-base-patch32") + embedder = descriptor.instantiate() + + if torch.cuda.is_available(): + expected_device = "cuda" + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + expected_device = "mps" + else: + expected_device = "cpu" + + embedder_device = next(embedder.model.parameters()).device + assert expected_device in str(embedder_device) + + +def test_transformers_pil_dependency_check(): + """Test that the embedder properly checks for PIL availability.""" + provider = TransformersProvider() + descriptor = provider.get_image_embedder("openai/clip-vit-base-patch32") + + # This should work when PIL is available (which it should be due to pytest.importorskip). + embedder = descriptor.instantiate() + assert embedder is not None + + +def test_transformers_pil_dependency_check_failure(): + """Test that the embedder fails fast when PIL is not available.""" + from unittest.mock import patch + + with patch("daft.dependencies.pil_image.module_available", return_value=False): + provider = TransformersProvider() + + with pytest.raises(ImportError, match="Pillow is required for image processing but not available"): + provider.get_image_embedder("openai/clip-vit-base-patch32") + + with pytest.raises(ImportError, match="Pillow is required for image processing but not available"): + test_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) + ( + daft.from_pydict({"image": [test_image]}) + .select(daft.col("image").cast(daft.DataType.image())) + .select(embed_image(daft.col("image"))) + ) From 7b2a06fa7750eb2383b0cea55cbb7c5d29f89b0b Mon Sep 17 00:00:00 2001 From: desmondcheongzx Date: Sun, 31 Aug 2025 14:20:07 -0700 Subject: [PATCH 03/11] address --- daft/ai/_expressions.py | 2 +- daft/ai/provider.py | 2 +- daft/ai/transformers/image_embedder.py | 4 +--- tests/ai/test_transformers.py | 4 ++-- 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/daft/ai/_expressions.py b/daft/ai/_expressions.py index 558f5cb0a4..53059c6dfa 100644 --- a/daft/ai/_expressions.py +++ b/daft/ai/_expressions.py @@ -22,7 +22,7 @@ def __call__(self, text_series: Series) -> list[Embedding]: class _ImageEmbedderExpression: - """Function expression implementation for a ImageEmbedder protocol.""" + """Function expression implementation for an ImageEmbedder protocol.""" image_embedder: ImageEmbedder diff --git a/daft/ai/provider.py b/daft/ai/provider.py index 41a1eb045b..6344d63d65 100644 --- a/daft/ai/provider.py +++ b/daft/ai/provider.py @@ -78,5 +78,5 @@ def get_text_embedder(self, model: str | None = None, **options: Any) -> TextEmb @abstractmethod def get_image_embedder(self, model: str | None = None, **options: Any) -> ImageEmbedderDescriptor: - """Returns a ImageEmbedderDescriptor for this provider.""" + """Returns an ImageEmbedderDescriptor for this provider.""" ... diff --git a/daft/ai/transformers/image_embedder.py b/daft/ai/transformers/image_embedder.py index 34dfea46da..0d13124fde 100644 --- a/daft/ai/transformers/image_embedder.py +++ b/daft/ai/transformers/image_embedder.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any import torch -from transformers import AutoModel, AutoProcessor +from transformers import AutoConfig, AutoModel, AutoProcessor from daft import DataType from daft.ai.protocols import ImageEmbedder, ImageEmbedderDescriptor @@ -31,8 +31,6 @@ def get_options(self) -> Options: return self.options def get_dimensions(self) -> EmbeddingDimensions: - from transformers import AutoConfig - config = AutoConfig.from_pretrained(self.model, trust_remote_code=True) # For CLIP models, the image embedding dimension is typically in projection_dim or hidden_size. embedding_size = getattr(config, "projection_dim", getattr(config, "hidden_size", 512)) diff --git a/tests/ai/test_transformers.py b/tests/ai/test_transformers.py index 169e76db2e..30f2bbf736 100644 --- a/tests/ai/test_transformers.py +++ b/tests/ai/test_transformers.py @@ -6,6 +6,8 @@ pytest.importorskip("torch") pytest.importorskip("PIL") +from unittest.mock import patch + import numpy as np import torch @@ -92,8 +94,6 @@ def test_transformers_pil_dependency_check(): def test_transformers_pil_dependency_check_failure(): """Test that the embedder fails fast when PIL is not available.""" - from unittest.mock import patch - with patch("daft.dependencies.pil_image.module_available", return_value=False): provider = TransformersProvider() From 7bae3036f9694c08baa7c8a1274c3578637d9d42 Mon Sep 17 00:00:00 2001 From: desmondcheongzx Date: Sun, 31 Aug 2025 14:21:14 -0700 Subject: [PATCH 04/11] fix mypy --- daft/ai/protocols.py | 4 ++-- daft/ai/transformers/image_embedder.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/daft/ai/protocols.py b/daft/ai/protocols.py index af48761e98..91dd28095d 100644 --- a/daft/ai/protocols.py +++ b/daft/ai/protocols.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import abstractmethod -from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable from daft.ai.typing import Descriptor @@ -31,7 +31,7 @@ def get_dimensions(self) -> EmbeddingDimensions: class ImageEmbedder(Protocol): """Protocol for image embedding implementations.""" - def embed_image(self, image: list[np.ndarray]) -> list[Embedding]: + def embed_image(self, image: list[np.ndarray[Any, Any]]) -> list[Embedding]: """Embeds a batch of images into an embedding vector.""" ... diff --git a/daft/ai/transformers/image_embedder.py b/daft/ai/transformers/image_embedder.py index 0d13124fde..485fb53b79 100644 --- a/daft/ai/transformers/image_embedder.py +++ b/daft/ai/transformers/image_embedder.py @@ -56,7 +56,7 @@ def __init__(self, model_name_or_path: str, **options: Any): self.processor = AutoProcessor.from_pretrained(model_name_or_path, trust_remote_code=True, use_fast=True) self.options = options - def embed_image(self, images: list[np.ndarray]) -> list[Embedding]: + def embed_image(self, images: list[np.ndarray[Any, Any]]) -> list[Embedding]: # TODO(desmond): There's potential for image decoding and processing on the GPU with greater # performance. Methods differ a little between different models, so let's do it later. pil_images = [pil_image.fromarray(image) for image in images] From a896aeef9b464b0be1f26b272dff139a13ee12ef Mon Sep 17 00:00:00 2001 From: desmondcheongzx Date: Sun, 31 Aug 2025 14:46:19 -0700 Subject: [PATCH 05/11] ease up on ci --- tests/ai/test_sentence_transformers.py | 7 ++++--- tests/ai/test_transformers.py | 11 ++++++----- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/ai/test_sentence_transformers.py b/tests/ai/test_sentence_transformers.py index 8b297ee4e9..6a8f5ef6a8 100644 --- a/tests/ai/test_sentence_transformers.py +++ b/tests/ai/test_sentence_transformers.py @@ -11,6 +11,7 @@ from daft.ai.sentence_transformers import SentenceTransformersProvider from daft.ai.typing import EmbeddingDimensions from daft.datatype import DataType +from tests.benchmarks.conftest import IS_CI def test_sentence_transformers_text_embedder_default(): @@ -23,7 +24,7 @@ def test_sentence_transformers_text_embedder_default(): @pytest.mark.parametrize( - "model_name, dimensions, run_model", + "model_name, dimensions, run_model_in_ci", [ ("sentence-transformers/all-MiniLM-L6-v2", 384, True), ("sentence-transformers/all-mpnet-base-v2", 768, True), @@ -33,7 +34,7 @@ def test_sentence_transformers_text_embedder_default(): ("BAAI/bge-base-en-v1.5", 768, True), ], ) -def test_sentence_transformers_text_embedder_other(model_name, dimensions, run_model): +def test_sentence_transformers_text_embedder_other(model_name, dimensions, run_model_in_ci): mock_options = {"arg1": "val1", "arg2": "val2"} provider = SentenceTransformersProvider() @@ -43,7 +44,7 @@ def test_sentence_transformers_text_embedder_other(model_name, dimensions, run_m assert descriptor.get_model() == model_name assert descriptor.get_options() == mock_options - if run_model: + if not IS_CI or run_model_in_ci: embedder = descriptor.instantiate() true_dimensions = embedder.model.get_sentence_embedding_dimension() diff --git a/tests/ai/test_transformers.py b/tests/ai/test_transformers.py index 30f2bbf736..a0d0c7302b 100644 --- a/tests/ai/test_transformers.py +++ b/tests/ai/test_transformers.py @@ -17,6 +17,7 @@ from daft.ai.typing import EmbeddingDimensions from daft.datatype import DataType from daft.functions.ai import embed_image +from tests.benchmarks.conftest import IS_CI def test_transformers_image_embedder_default(): @@ -29,15 +30,15 @@ def test_transformers_image_embedder_default(): @pytest.mark.parametrize( - "model_name, dimensions, run_model", + "model_name, dimensions, run_model_in_ci", [ ("openai/clip-vit-base-patch32", 512, True), - ("openai/clip-vit-large-patch14", 768, True), + ("openai/clip-vit-large-patch14", 768, False), ("openai/clip-vit-base-patch16", 512, True), - ("openai/clip-vit-large-patch14-336", 768, True), + ("openai/clip-vit-large-patch14-336", 768, False), ], ) -def test_transformers_image_embedder_other(model_name, dimensions, run_model): +def test_transformers_image_embedder_other(model_name, dimensions, run_model_in_ci): mock_options = {"arg1": "val1", "arg2": "val2"} provider = TransformersProvider() @@ -47,7 +48,7 @@ def test_transformers_image_embedder_other(model_name, dimensions, run_model): assert descriptor.get_model() == model_name assert descriptor.get_options() == mock_options - if run_model: + if not IS_CI or run_model_in_ci: embedder = descriptor.instantiate() # Test with a variety of image sizes and shapes that should be preprocessed correctly. From e50ecd3e15c6a4f16491969e370803c7c60ef98a Mon Sep 17 00:00:00 2001 From: desmondcheongzx Date: Sun, 31 Aug 2025 16:08:14 -0700 Subject: [PATCH 06/11] fix ci --- docs/install.md | 14 +++++++++++--- pyproject.toml | 5 +++-- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/docs/install.md b/docs/install.md index d5852d15f2..95ccf9d83b 100644 --- a/docs/install.md +++ b/docs/install.md @@ -41,6 +41,14 @@ Depending on your use case, you may need to install Daft with additional depende + +