Skip this notebook execution in CI because it requires AWS credentials for presigned URL generation - Daft

python

!pip install daft deltalake<0.17

python

CI = False

python

# Skip this notebook execution in CI because it requires AWS credentials for presigned URL generation
if CI:
    import sys

    sys.exit()

Multimodal data lake annotation and indexing

Let's go from: Images in an S3 Bucket

To: Multimodal Data Lake where we can run queries efficiently to power analytics, retrieval and more!

python

import daft

IO_CONFIG = daft.io.IOConfig(s3=daft.io.S3Config(anonymous=True))  # Use anonymous S3 access

daft.set_planning_config(default_io_config=IO_CONFIG)

python

df = daft.from_glob_path(
    "s3://daft-public-data/open-images/validation-images/*",
)
df.show()

Working with URLs in Daft is really easy and efficient

URLs are extremely common when working with multimodal data, most commonly as a https:// URL or s3:// object store URL
Daft runs URL downloads using async Rust kernels, saturating your machine's network bandwidth even for millions of small files (see: demo at PyData Global 2023)

python

df = df.with_column("image_bytes", df["path"].download())
df.show()

Reading Images

Daft makes working with opaque file formats/encodings easy

Native type available for images and tensors
Support for arbitrary Python objects in columns so you can use all your favorite Python libraries as well for datatypes not yet supported by Daft (e.g. video, audio, PDFs)

python

df = df.with_column("image", df["image_bytes"].decode_image())
df.show()

Thumbnail creation

Easily create thumbnails for your image using the resize(...) Daft expression.

python

df = df.with_column("image_thumbnail", df["image"].resize(32, 32))
df.show()

Running multimodal LLMs

Since we are running on just our laptop, we will be offloading our "heavy compute" (running the GPT-4o model on our image) to the OpenAI API.

If instead we wanted to run our own models or algorithms, Daft also lets us run on GPUs by adding gpus=1 to the @daft.cls (or @daft.func) decorator — see the GPU guide.

python

import json
import os

import aiohttp
import boto3

DEFAULT_PROMPT = "What’s in this image?"
api_key = os.getenv("OPENAI_API_KEY")
if api_key is None:
    raise RuntimeError("Please specify your OpenAI API key as the environment variable `OPENAI_API_KEY`.")

headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}


# Row-wise @daft.func: Daft turns the Python return type hint into the Daft dtype,
# and calls us once per row. Daft handles batching and Series construction.
@daft.func
def generate_presigned_url(s3_url: str, expires_in: int = 3600) -> str:
    """Generate a presigned Amazon S3 URL."""
    s3_client = boto3.client("s3")
    bucket, key = s3_url.strip("s3://").split("/", 1)
    return s3_client.generate_presigned_url(
        ClientMethod="get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=expires_in
    )


# Async row-wise @daft.func: per-row HTTP request runs concurrently up to
# max_concurrency coroutines. Daft awaits them for us.
@daft.func(max_concurrency=16)
async def run_gpt4o_on_urls(image_url: str, prompt: str = DEFAULT_PROMPT) -> str:
    """Run the gpt-4o LLM by making an API call to OpenAI."""
    payload = {
        "model": "gpt-4o",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": image_url}},
                ],
            }
        ],
        "max_tokens": 300,
    }

    async with aiohttp.ClientSession(headers=headers) as session:
        async with session.post("https://api.openai.com/v1/chat/completions", json=payload) as response:
            return json.dumps(await response.json())

python

# Generate temporary URLs with a short expiration time
df = df.with_column("image_urls", generate_presigned_url(df["path"]))

# Make remote API calls to OpenAI endpoint
df = df.with_column("gpt_results", run_gpt4o_on_urls(df["image_urls"], prompt="What’s in this image?"))

# Parse JSON outputs from OpenAI endpoint
df = df.with_column("description", df["gpt_results"].json.query(".choices[0].message.content"))

df.show(3)

python

df = df.select(
    # Larger multimodal data (such as large images or documents) can be written as URLs
    "path",
    # Small multimodal data (such as thumbnails or full-form text) can be written inline
    df["image_thumbnail"].encode_image("JPEG"),
    # Metadata such as size in bytes and descriptions should be stored as per normal
    "size",
    "description",
)

python

df

python

# Limit to running just 8 rows to save your OpenAI bill...
df = df.limit(8)

df.write_delta("my_table.delta_lake")

Now we have our "Multimodal Data Lake"!

Thumbnails readily available for visualization
URLs available for access to the raw data
Extracted metadata (description) available for querying

python

read_df = daft.read_deltalake("my_table.delta_lake")
read_df

python

read_df = read_df.with_column("image_thumbnail", daft.col("image_thumbnail").decode_image()).where(
    read_df["description"].contains("dog")
)

python

read_df.collect()