Distilabel

Distilabelは、検証済みの研究論文に基づいた高速・信頼性・スケーラビリティを備えたパイプラインを必要とするエンジニア向けの、合成データおよびAIフィードバックのためのフレームワークです。

Distilabelは、従来の予測型NLP（分類、抽出など）から、生成型および大規模言語モデルのシナリオ（指示追従、対話生成、ジャッジングなど）まで、幅広いプロジェクトにおける合成データおよびAIフィードバックの生成に利用できます。Distilabelのプログラムによるアプローチにより、スケーラブルなデータ生成およびAIフィードバックのパイプラインを構築できます。distilabelの目標は、AIフィードバックを用いた生成およびジャッジングに関する検証済みの研究手法に基づき、高品質で多様なデータセットを迅速に生成することで、AI開発を加速させることです。

distilabelを使って、人々はどのようなものを作っているのでしょうか？

Argillaコミュニティは、distilabelを使用して素晴らしいデータセットとモデルを作成しています。

• 1M OpenHermesPreferenceは、teknium/OpenHermes-2.5 LLMを使用して生成された約100万件のAI好みをまとめたデータセットです。distilabelを使用してデータセットの開発をスケーリングし、規模を拡大する方法の優れた例です。
• distilabeled Intel Orca DPO データセットは、改善された OpenHermes モデルのファインチューニングに使用されました。このデータセットは、Argilla による人間のキュレーションと distilabel からの AI フィードバックを組み合わせることで構築され、Intel Orca データセットの改善版となり、元のデータセットでファインチューニングされたモデルを上回る性能を発揮しました。
• haiku DPOデータは、特定のタスク用の合成データセットを誰でも作成できることを示す例であり、キュレーションと評価を経て、カスタムLLMのファインチューニングに使用できます。

前提条件

まず、Hugging Face アカウントでログインします:

hf auth login

distilabelがインストールされていることを確認してください:

pip install -U distilabel[vllm]

Distilabel パイプライン

Distilabelのパイプラインは、相互に接続された任意の数のステップまたはタスクで構築できます。1つのステップまたはタスクの出力は、別のステップまたはタスクの入力として渡されます。複数のステップを連鎖させることで、LLMを用いた複雑なデータ処理および生成パイプラインを構築できます。各ステップの入力はデータのバッチであり、辞書のリストを含みます。各辞書はデータセットの1行を表し、キーは列名です。Hugging Faceハブからデータを取得し、ハブにデータを出力するために、datasets.DatasetDictの抽象化としてDistisetクラスを定義しました。

Distisetをデータセットオブジェクトとして

distilabel の Pipeline は、Distiset と呼ばれる Hugging Face の datasets.DatasetDict の特殊な型を返します。

Pipeline は、Distiset として複数のサブセットを出力できます。Distiset は、サブセットごとに1つのエントリを持つ辞書のようなオブジェクトです。その後、Distiset は Hugging Face Hub にシームレスにプッシュでき、すべてのサブセットが同じリポジトリに格納されます。

HubからDistisetにデータを読み込む

Hub からデータを読み込む例として、Prometheus 2 paperを再現し、distilabel に実装されている PrometheusEval タスクを使用します。Prometheus 2 と Prometheuseval タスクは、それぞれ直接評価タスクとペアワイズランキングタスクです。前者は、参照回答の有無に関わらず、与えられた指示に対する単一の孤立した応答の品質を評価し、後者は、参照回答の有無に関わらず、与えられた指示に対する 1 つの応答を別の応答と比較して品質を評価します。これらのタスクは、Hugging Face の H4 チームによって作成され、Hub から読み込まれる HuggingFaceH4/instruction-dataset というデータセットに対して使用します。

from distilabel.llms import vLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import KeepColumns, LoadDataFromHub
from distilabel.steps.tasks import PrometheusEval

if __name__ == "__main__":
    with Pipeline(name="prometheus") as pipeline:
        load_dataset = LoadDataFromHub(
            name="load_dataset",
            repo_id="HuggingFaceH4/instruction-dataset",
            split="test",
            output_mappings={"prompt": "instruction", "completion": "generation"},
        )

        task = PrometheusEval(
            name="task",
            llm=vLLM(
                model="prometheus-eval/prometheus-7b-v2.0",
                chat_template="[INST] {{ messages[0]['content'] }}\n{{ messages[1]['content'] }}[/INST]",
            ),
            mode="absolute",
            rubric="factual-validity",
            reference=False,
            num_generations=1,
            group_generations=False,
        )

        keep_columns = KeepColumns(
            name="keep_columns",
            columns=["instruction", "generation", "feedback", "result", "model_name"],
        )

        load_dataset >> task >> keep_columns

次に、パイプラインを起動し、データをDistisetオブジェクトに保存できるようにするため、ランタイムパラメータを指定してpipeline.runを呼び出す必要があります。

distiset = pipeline.run(
    parameters={
        task.name: {
            "llm": {
                "generation_kwargs": {
                    "max_new_tokens": 1024,
                    "temperature": 0.7,
                },
            },
        },
    },
)

distilabelのDistisetをHubにプッシュする

DistisetをHugging Faceのリポジトリにプッシュし、各サブセットが異なる設定に対応するようにします。

distiset.push_to_hub(
    "my-org/my-dataset",
    commit_message="Initial commit",
    private=False,
    token=os.getenv("HF_TOKEN"),
)

📚 リソース

• 🚀 Distilabel ドキュメント
• 🚀 Distilabel ドキュメント - distiset
• 🚀 Distilabel ドキュメント - prometheus
• 🆕 distilabel の紹介