DataChain 🔨 Wrangle unstructured AI data at scale

2024年MLOpsアドベントカレンダー 12/22の記事です！

生成AIの開発と利用がどんどん興隆している昨今では、画像や音声ファイル、PDF、テキスト文書など膨大な非構造化データを扱う必要があります。筆者も生成AIのモデル開発を実践している会社で日々データエンジニアリングを担っていますが、多種多様で膨大なマルチモーダルデータを効率的に管理、変換、分析することは重大な課題となっています。

DataChainは非構造化データのために設計されたPythonベースのAIデータウェアハウスであるDataChainについて、その目的、用途、特徴、使用例、およびメリットを詳しく解説します。

GitHubレポジトリ：https://github.com/iterative/datachain
Document：https://docs.datachain.ai/

なお、DataChainを開発しているIterativeはDVCを開発していることでも有名な会社です。

DataChainとは？

DataChainは、非構造化データをスムーズに変換、分析、バージョン管理できるETLツールです。Amazon S3、GCS、Azure Blobなどの外部ストレージやHuggingFaceのようなモデルレポジトリと統合し、効率的にデータ処理を実現します。内部データベースでメタデータを管理することで、大規模なデータセットに対する複雑なクエリと操作を容易にします。

DataChainは以下のようなコアモジュールを提供しています：

DataChain: データ操作の中心的な機能を提供。
DataType: 各種データ型を扱うためのツール。
File: ファイル操作に特化したメソッド群。
UDF（ユーザー定義関数）: データ変換や操作をカスタマイズするためのインターフェース。
Torchサポート: PyTorchを使用した機械学習モデルとの連携。
SQLっぽいインタフェース: 高度なクエリ操作を可能にするSQLサポート。

なお、DataChainのソースコードも上記のような構成で整理されており、だいぶ読みやすくなっています。

DataChainの価値

DataChainが提供する価値は次の通りです。

非構造化データ処理の効率化: 画像、動画、音声、PDFなどの非構造化データの取り扱いを簡素化。
データ変換の促進: Python的なワークフローを用いて、データを強化、変換、分析。
AIモデルトレーニングの最適化: 機械学習および深層学習モデルのためのデータ準備と管理の堅牢なフレームワークを提供。
コラボレーションと再現性の向上: データセットのバージョン管理と効率的な共有を可能にするツールを提供。

生成AIのモデル開発におけるデータエンジニアリングの課題

生成AIのモデル開発では、膨大かつ多様なデータセットを扱う必要があります。このプロセスで以下のような課題が発生します：

データの品質管理:
- 収集したデータにノイズや重複が含まれることが多い。
- 不適切なデータや有害なデータのフィルタリングが必要。
- 高品質なデータセットを維持するには、ルールベースやモデルベースのフィルタリングが必須。
データの多様性:
- モデルの汎用性を向上させるために、多様なドメインや形式のデータが求められる。
- 特定のタスクや言語に最適化されたデータセットを作成する必要がある。
データのスケーラビリティ:
- 数百万から数十億件のデータを処理する際、ストレージや計算リソースの効率的な利用が課題。
- ストレージや計算負荷を抑えるための最適化が求められる。
データの変換と統一:
- 異なるフォーマットや構造のデータを統一的に処理する必要がある。
- メタデータの生成や付加による効率的な検索と管理が重要。
PIIや倫理的懸念の対応:
- 個人情報やセンシティブなデータを削除またはマスクする必要がある。

たとえば私の経験で言うと、動画データを大量に収集して動画生成AIを開発するケースを考えます。データソース次第ですが、収集したデータの品質が高いとは限らないことは多々あります。現時点（2024年末）での動画生成AI（Text-to-video）の主流は数秒程度の短い動画の生成ですが、学習データも同様に数秒程度の短いクリップ動画とキャプションテキストが必要になります。もちろん世の中の動画データのほとんどは数秒どころか数分、数時間で作られており、学習データはその中から切り出して来る必要があります。さらにはクリップ動画の重複排除、モーションスコアリング、Aesthetic scoring、PIIや不適切動画の排除、不要なテキスト（字幕等）の排除、キャプションの作成といった工程が必要になります。学習データを用意するだけで一大プロジェクトになり、複雑なデータパイプラインを開発してインフラを管理する必要があります。

こうした課題に対処するとき、共通化されたインタフェースとETL仕様を決めることはチームで効率的に開発するために必須の活動であり、DataChainのようにデータ処理とインタフェース、データタイプが揃ったライブラリはとてもありがたいです。

Meta Movie Gen
Open-Sora Plan: Open-Source Large Video Generation Model
OpenVid-1M: A Large-Scale High-Quality Dataset for Text-to-video Generation

DataChainが提供する解決策

DataChainは、ETLプロセスを効率化するための以下のような機能を提供します。

柔軟なデータ抽出:
- DataChainのfrom_storageメソッドを使用して、複数のストレージから簡単にデータを抽出可能。
- S3、GCP、Azure、HuggingFaceなどの主要ストレージプラットフォームに対応。
高度なデータ変換とクレンジング:
- PythonベースのUDFを利用して、カスタマイズ可能なデータ変換を実現。
- データの重複排除や、ノイズ削減のための高度なアルゴリズムをサポート。
- メタデータの自動生成機能を備え、効率的なデータ管理を実現。
効率的なデータロード:
- フィルタリングや変換後のデータをスムーズにモデル学習環境にロード。
- PyTorchと統合してモデルトレーニングにシームレスに接続可能。PyTorch data loaderにそのまま繋げられるのは便利。

生成AIにおける実践例

生成AIプロジェクトにおいて、DataChainは以下のような場面で具体的に活用されています。

画像データの評価:
学習済モデルを用いて画像データセットを評価。
LLMを用いたデータの評価:
LLMのAPIを用いてテキストデータを評価。
画像に対するキャプションテキストの作成:
マルチモーダルモデルを用いて入力画像に対する説明文を生成。
PyTorch data loaderとの統合:
PyTorch data loaderにDataChain統合。

利用例とコードサンプル

1. 画像データセットの評価

学習済モデルを用いて画像データセットを評価します。

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")

2. LLMを用いたデータの評価

LLMの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")

3. 画像に対するキャプションテキストの作成

マルチモーダルモデルを用いて入力画像に対する説明文を生成します。

from datachain.lib.dc import Column, DataChain, File 
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration 

images = DataChain.from_storage("gs://datachain-demo/newyorker_caption_contest/images", type="image")

model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma-3b-mix-224")
processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")

def process(file: File) -> str:
  image=file.read().convert("RGB")
  inputs = processor(text="caption", images=image, return_tensors="pt")
  generate_ids = model.generate(**inputs, max_new_tokens=100)
  return processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]

chain = (
      images.limit(5)
      .settings(cache=True)
      .map(scene=lambda file: process(file), output = str)
      .save()
)

4. PyTorch data loaderとの統合

PyTorch data loaderにDataChain統合します。

ds = train.select("file", "caption_choices", "label_ind").to_pytorch(
    transform=preprocess,
    tokenizer=clip.tokenize,
)

loader = DataLoader(ds, batch_size=2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
train(loader, model, optimizer)

その他のTutorialも用意されています。

結論

DataChainは非構造化データを管理、変換、分析するためのツールであり、特にAIや機械学習プロジェクトにおいてその真価を発揮するでしょう。多様なデータタイプやストレージとの統合が可能なDataChainは、データエンジニアリングからAIモデルのトレーニングまで、幅広い用途に対応しています。また、生成AI開発における具体的な課題を解決し、エンジニアが直面する複雑なデータ処理フローを簡素化する強力な機能を提供します。

生成AIのための非構造化マルチモーダルデータのETLやデータエンジニアリングはまだまだ発展途上で、注目されることも少ないですが、データはAI開発全般の最重要要素であるので、DataChain含めて技術やライブラリが広まっていくことを期待しています。