AI 🔗 DataChain

DataChain is an open-source Python data processing library for wrangling unstructured AI data at scale.

It enables batch LLM API calls and local language and vision AI model inferences to run in parallel over many samples as chained operations resolving to table-like datasets. These datasets can be saved, versioned, and sent directly to PyTorch and TensorFlow for training. DataChain employs rigorous Pydantic data structures, promoting better data processing practices and enabling vectorized analytical operations normally found in databases.

The DataChain fills the gap between dataframe libraries, data warehouses, and Python-based multimodal AI applications. Our primary use cases include massive data curation, LLM analytics and validation, batch image segmentation and pose detection, GenAI data alignment, etc.

$ pip install data-chain

Basic operation

DataChain is built by composing wrangling operations.

For example, it can be instructed to read files from the cloud, map them onto a modern AI service returning a Python object, parallelize API calls, save the result as a dataset, and export a column:

import os
import data_chain as dc

from anthropic.types.message import Message
ClaudeModel = dc.pydantic_to_feature(Message)
PROMPT = "summarize this book in less than 200 words"
service = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
source = "gs://datachain-demo/mybooks/"

chain = dc.DataChain(source)                          \
              .filter(File.name.glob("*.txt"))        \
              .settings(parallel=4)                   \
              .map(                                   \
                             claude = lambda file:                                         \
                                       ClaudeModel(**service.messages.create(                        \
                                         model="claude-3-haiku-20240307",         \
                                         system=PROMPT,                           \
                                         messages=[{"role": "user",               \
                                                    "content": file.get_value()}] \
                                                               ),  \
                                                       ).model_dump()  \
                                                       )               \
                                                       .save("mydataset")

dc.DataChain("mydataset").export("./", "claude.response") # export summaries

Dataset persistence

In the example above, the chain resolves to a saved dataset “mydataset”. DataChain datasets are immutable and versioned. A saved dataset version can be used as a data source:

ds = dc.DataChain("mydataset", version = 1)

Note that DataChain represents file samples as pointers into their respective storage locations. This means a newly created dataset version does not duplicate files in storage, and storage remains the single source of truth for the original samples

Vectorized analytics

Since datasets are internally represented as tables, analytical queries can be vectorized:

rate = ds.filter(chain.response == "Success").count() / chain.count() # ??
print(f"API class success rate: {100*rate:.2f}%")
>> 74.68%

price_input = 0.25
price_output = 1.25
price=(ds.sum(C.claude.usage.input_tokens)*price_input \
       + ds.sum(C.claude.usage.output_tokens)*price_output)/1_000_000
print(f"Cost of API calls: ${price:.2f}")
>> Cost of API calls: $1.42

Importing metadata

It is common for AI data to come together with metadata (annotations, classes, etc). DataChain understands many metadata formats, and can connect data samples in storage with external metadata (e.g. CSV columns) to form a single dataset:

from dc import parse_csv

files = dc.DataChain("gs://datachain-demo/myimages/")
metadata = dc.DataChain("gs://datachain-demo/myimagesmetadata.csv") \
               .gen(meta=parse_csv)  # TBD, also dependent on dropping file
dataset = chain1.merge(chain2, on = "file.name", right_on="name"])

print(dataset.select("file.name", "class", "prob").limit(5).to_pandas())
....
....
....
....
....

Nested annotations (like JSON) can be unrolled into rows and columns in the way that best fits the application. For example, the MS COCO dataset includes JSON annotations detailing segmentations. To build a dataset consisting of all segmented objects in all COCO images:

image_files = dc.DataChain("gs://datachain-demo/coco/images/")
image_meta  = dc.DataChain("gs://datachain-demo/coco.json")  \
               .gen(meta=parse_json, key="images")       # list of images
images = image_files.merge(image_meta, on = "file.name", right_on="file_name")
objects_meta = dc.DataChain("gs://datachain-demo/coco.json") \
               .gen(meta=parse_json, key="annotations")  # annotated objects

objects = image.full_merge(objects_meta, on = "id", right_on = "image_id")

Generating metadata

A typical step in data curation is to create features from data samples for future selection. DataChain represents the newly created metadata as columns, which makes it easy to create new features and filter on them:

from fashion_clip.fashion_clip import FashionCLIP
from sqlalchemy import JSON
from tabulate import tabulate

from datachain.lib.param import Image
from datachain.query import C, DatasetQuery, udf


@udf(
    params=(Image(),),
    output={"fclip": JSON},
    method="fashion_clip",
    batch=10,
)
class MyFashionClip:
    def __init__(self):
        self.fclip = FashionCLIP("fashion-clip")

    def fashion_clip(self, inputs):
        embeddings = self.fclip.encode_images(
            [input[0] for input in inputs], batch_size=1
        )
        return [(json.dumps(emb),) for emb in embeddings.tolist()]

chain = dc.DataChain("gs://datachain-demo/zalando/images/").filter(
        C.name.glob("*.jpg")
    ).limit(5).add_signals(MyFashionClip).save("zalando_hd_emb")

test_image = "cs://datachain-demo/zalando/test/banner.jpg"
test_embedding = MyFashionClip.fashion_clip.encode_images(Image(test_image))

best_matches = chain.filter(similarity_search(test_embeding)).limit(5)

print best_matches.to_result()

Delta updates

DataChain is capable of “delta updates” – that is, batch-processing only the newly added data samples. For example, let us copy some images into a local folder and run a chain to generate captions with a locally served captioning model from HuggingFace:

> mkdir demo-images/
> datachain cp gs://datachain-demo/images/ /tmp/demo-images

import torch

from datachain.lib.hf_image_to_text import LLaVAdescribe
from datachain.query import C, DatasetQuery

source = "/tmp/demo-images"

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

if __name__ == "__main__":
    results = (
        DatasetQuery(
            source,
            anon=True,
        )
        .filter(C.name.glob("*.jpg"))
        .add_signals(
            LLaVAdescribe(
                device=device,
                model=model,
            ),
            parallel=False,
        )
        .save("annotated-images")
    )

Now let us add few more more images to the same folder:

> datachain cp gs://datachain-demo/extra-images/ /tmp/demo-images

and calculate updates only for the delta:

processed = dc.DataChain("annotated-images")
delta = dc.dataChain("/tmp/demo-images").subtract(processed)

Passing data to training

Datasets can be exported to CSV or webdataset formats. However, a much better way to pass data to training which avoids data copies and re-sharding is to wrap a DataChain dataset into a PyTorch class, and let the library take care of file downloads and caching under the hood:

ds = dc.DataChain("gs://datachain-demo/name-labeled/images/")
               .filter(C.name.glob("*.jpg"))
               .map(lambda name: (name[:3],), output={"label": str}, parallel=4)
    )

train_loader = DataLoader(
        ds.to_pytorch(
            ImageReader(),
            LabelReader("label", classes=CLASSES),
            transform=transform,
        ),
        batch_size=16,
        parallel=2,
    )

💻  More examples

• Curating images to train a custom CLIP model without re-sharding the Webdataset files

• Batch-transforming and indexing images to create a searchable merchandise catalog

• Evaluating an LLM application at scale

• Ranking the LLM retrieval strategies

• Delta updates in batch processing

Contributions

Contributions are very welcome. To learn more, see the Contributor Guide.

License

Distributed under the terms of the Apache 2.0 license, DataChain is free and open source software.

Issues

If you encounter any problems, please file an issue along with a detailed description.