DataChain: 现代化的AI数据处理框架

DataChain:为AI而生的数据处理利器

在人工智能和机器学习快速发展的今天,如何高效地处理和管理大规模非结构化数据成为了一个关键挑战。DataChain应运而生,它是一个专为AI设计的现代Python数据框架库,旨在帮助数据科学家和机器学习工程师更好地组织和处理复杂的数据集。

DataChain的核心特性

DataChain具有几个突出的特性,使其在AI数据处理领域脱颖而出:

以存储为真实来源
- 无需冗余复制即可处理来自S3、GCP、Azure和本地文件系统的非结构化数据
- 支持多模态数据:图像、视频、文本、PDF、JSON、CSV、Parquet等
- 将文件和元数据统一到持久化、版本化的列式数据集中
Python友好的数据管道
- 直接操作Python对象和对象字段
- 内置并行化和内存外计算,无需SQL或Spark
数据丰富和处理
- 使用本地AI模型和LLM API生成元数据
- 通过元数据进行过滤、连接和分组,支持向量嵌入搜索
- 将数据集传递给PyTorch和TensorFlow,或导出回存储
高效性能
- 支持并行化、内存外工作负载和数据缓存
- 对Python对象字段进行向量化操作:求和、计数、平均等
- 优化的向量搜索

DataChain架构图

快速上手DataChain

让我们通过几个实际的例子来了解DataChain的强大功能:

使用JSON元数据选择文件

假设我们有一个存储库,包含猫和狗的图像(如dog.1048.jpg, cat.1009.jpg),每个图像都有一个匹配的JSON文件(如cat.1009.json)包含地面真实标签和模型推理结果:

{
    "class": "cat", "id": "1009", "num_annotators": 8,
    "inference": {"class": "dog", "confidence": 0.68}
}

使用DataChain,我们可以轻松地只下载那些被高置信度推断为"猫"的图像:

from datachain import Column, DataChain

meta = DataChain.from_json("gs://datachain-demo/dogs-and-cats/*json", object_name="meta")
images = DataChain.from_storage("gs://datachain-demo/dogs-and-cats/*jpg")

images_id = images.map(id=lambda file: file.path.split('.')[-2])
annotated = images_id.merge(meta, on="id", right_on="meta.id")

likely_cats = annotated.filter((Column("meta.inference.confidence") > 0.93) & 
                               (Column("meta.inference.class_") == "cat"))
likely_cats.export_files("high-confidence-cats/", signal="file")

这个例子展示了DataChain如何轻松地处理复杂的数据选择任务,结合文件存储和元数据信息。

使用本地AI模型进行数据整理

DataChain还支持使用本地AI模型进行批量推理。以下是使用简单情感分析模型处理文本文件的例子:

from transformers import pipeline
from datachain import DataChain, Column

classifier = pipeline("sentiment-analysis", device="cpu",
                model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")

def is_positive_dialogue_ending(file) -> bool:
    dialogue_ending = file.read()[-512:]
    return classifier(dialogue_ending)[0]["label"] == "POSITIVE"

chain = (
   DataChain.from_storage("gs://datachain-demo/chatbot-KiT/",
                          object_name="file", type="text")
   .settings(parallel=8, cache=True)
   .map(is_positive=is_positive_dialogue_ending)
   .save("file_response")
)

positive_chain = chain.filter(Column("is_positive") == True)
positive_chain.export_files("./output")

print(f"{positive_chain.count()} files were exported")

这个例子展示了如何使用DataChain结合本地AI模型进行大规模数据处理和筛选。

LLM评估聊天机器人

DataChain还可以与大型语言模型(LLM)API集成,用于更复杂的数据处理任务。以下是使用Mistral API评估聊天机器人对话的例子:

from mistralai import Mistral
from datachain import File, DataChain, Column

PROMPT = "Was this dialog successful? Answer in a single word: Success or Failure."

def eval_dialogue(file: File) -> bool:
     client = Mistral()
     response = client.chat.complete(
         model="open-mixtral-8x22b",
         messages=[{"role": "system", "content": PROMPT},
                   {"role": "user", "content": file.read()}])
     result = response.choices[0].message.content
     return result.lower().startswith("success")

chain = (
   DataChain.from_storage("gs://datachain-demo/chatbot-KiT/", object_name="file")
   .settings(parallel=4, cache=True)
   .map(is_success=eval_dialogue)
   .save("mistral_files")
)

successful_chain = chain.filter(Column("is_success") == True)
successful_chain.export_files("./output_mistral")

print(f"{successful_chain.count()} files were exported")

这个例子展示了DataChain如何与LLM API集成,实现复杂的数据评估任务。

序列化Python对象

DataChain的一个强大特性是能够序列化整个LLM响应到内部数据库:

from mistralai import Mistral
from mistralai.models import ChatCompletionResponse
from datachain import File, DataChain, Column

PROMPT = "Was this dialog successful? Answer in a single word: Success or Failure."

def eval_dialog(file: File) -> ChatCompletionResponse:
     client = MistralClient()
     return client.chat(
         model="open-mixtral-8x22b",
         messages=[{"role": "system", "content": PROMPT},
                   {"role": "user", "content": file.read()}])

chain = (
   DataChain.from_storage("gs://datachain-demo/chatbot-KiT/", object_name="file")
   .settings(parallel=4, cache=True)
   .map(response=eval_dialog)
   .map(status=lambda response: response.choices[0].message.content.lower()[:7])
   .save("response")
)

chain.select("file.name", "status", "response.usage").show(5)

success_rate = chain.filter(Column("status") == "success").count() / chain.count()
print(f"{100*success_rate:.1f}% dialogs were successful")

这个例子展示了如何使用DataChain序列化和分析复杂的LLM响应。