Not Diamond のカスタムルーティング

これはインタラクティブなノートブックです。ローカルで実行するか、以下のリンクから利用できます。

このノートブックでは、Weave と Not Diamond のカスタムルーティングを使用して、評価結果に基づき LLM プロンプトを最適なモデルにルーティングする方法を紹介します。このノートブックを最後まで進めると、コーディング用プロンプト向けのカスタムルーターをトレーニングし、それを使って新しいプロンプトのターゲットモデルを選択したうえで、単体で最も優れたモデルとパフォーマンスを比較評価できます。

プロンプトのルーティング

このセクションでは、複数の LLM を扱う際にカスタムルーティングが役立つ理由を説明します。複雑な LLM ワークフローを構築する際には、精度、コスト、呼び出しのレイテンシーに応じて、プロンプトを送るモデルを使い分ける必要が生じることがあります。 Not Diamond を使用すると、こうしたワークフロー内のプロンプトをニーズに最適なモデルへルーティングでき、モデルのコストを抑えつつ精度を最大限に高められます。どのようなデータ分布であっても、単一のモデルがあらゆるクエリで他のすべてのモデルを上回ることはまれです。複数のモデルを組み合わせ、どの LLM をいつ呼び出すべきかを学習する「メタモデル」を構築することで、個々のどのモデルよりも高いパフォーマンスを実現できるうえ、コストとレイテンシーも削減できます。

カスタムルーティング

プロンプト用のカスタムルーターをトレーニングするには、次の入力が必要です。

LLM プロンプトのセット: プロンプトは文字列である必要があります。また、アプリケーションで実際に使用されるプロンプトを代表するものにしてください。
LLM の応答: 各入力に対する候補 LLM の応答です。候補 LLM には、サポートされる LLM と独自のカスタムモデルの両方を含めることができます。
入力に対する候補 LLM の応答の評価スコア: スコアは数値で、ニーズに合った任意のメトリクスを使用できます。

これらを Not Diamond API に送信すると、ワークフローごとに調整されたカスタムルーターをトレーニングできます。

トレーニングデータを設定する

このセクションでは、カスタムルーターの学習に使用するトレーニングデータとテストデータを準備します。実際の運用では、独自の評価を使用してカスタムルーターをトレーニングします。ただし、このサンプルノートブックでは、HumanEval データセットに対する LLM の応答を使用して、コーディングタスク向けのカスタムルーターをトレーニングします。まず、この例のために用意されたデータセットをダウンロードし、LLM の応答をモデルごとに解析して EvaluationResults に変換します。このトレーニングデータとテストデータの分割は、以降のセクションでルーターをトレーニングし、サンプル外データでのパフォーマンスを評価する際に使用します。

!curl -L "https://drive.google.com/uc?export=download&id=1q1zNZHioy9B7M-WRjsJPkfvFosfaHX38" -o humaneval.csv
python
import random

import weave
from weave.flow.dataset import Dataset
from weave.flow.eval import EvaluationResults
from weave.integrations.notdiamond.util import get_model_evals

pct_train = 0.8
pct_test = 1 - pct_train

# 実際には、独自のデータセットで評価を構築し、
# `evaluation.get_eval_results(model)` を呼び出します
model_evals = get_model_evals("./humaneval.csv")
model_train = {}
model_test = {}
for model, evaluation_results in model_evals.items():
    n_results = len(evaluation_results.rows)
    all_idxs = list(range(n_results))
    train_idxs = random.sample(all_idxs, k=int(n_results * pct_train))
    test_idxs = [idx for idx in all_idxs if idx not in train_idxs]

    model_train[model] = EvaluationResults(
        rows=weave.Table([evaluation_results.rows[idx] for idx in train_idxs])
    )
    model_test[model] = Dataset(
        rows=weave.Table([evaluation_results.rows[idx] for idx in test_idxs])
    )
    print(
        f"Found {len(train_idxs)} train rows and {len(test_idxs)} test rows for {model}."
    )

カスタムルーターをトレーニングする

EvaluationResults の準備ができたら、それを Not Diamond に送信してカスタムルーターをトレーニングできます。アカウントを作成して APIキーを生成したうえで、以下のコードに APIキーを入力してください。APIキーはトレーニングリクエストの認証に使用されます。

import os

from weave.integrations.notdiamond.custom_router import train_router

api_key = os.getenv("NOTDIAMOND_API_KEY", "<YOUR_API_KEY>")

preference_id = train_router(
    model_evals=model_train,
    prompt_column="prompt",
    response_column="actual",
    language="en",
    maximize=True,
    api_key=api_key,
    # 初めてカスタムルーターをトレーニングする場合は、この行をコメントアウトしたままにします
    # 既存のカスタムルーターを上書きして再トレーニングする場合は、この行のコメントを解除します
    # preference_id=preference_id,
)

その後、Not Diamond アプリでカスタムルーターのトレーニングの進行状況を確認できます。

カスタムルーターのトレーニングが完了したら、そのルーターを使用してプロンプトをルーティングできます。

from notdiamond import NotDiamond

import weave

weave.init("notdiamond-quickstart")

llm_configs = [
    "anthropic/claude-3-5-sonnet-20240620",
    "openai/gpt-4o-2024-05-13",
    "google/gemini-1.5-pro-latest",
    "openai/gpt-4-turbo-2024-04-09",
    "anthropic/claude-3-opus-20240229",
]
client = NotDiamond(api_key=api_key, llm_configs=llm_configs)

new_prompt = (
    """
You are a helpful coding assistant. Using the provided function signature, write the implementation for the function
in Python. Write only the function. Do not include any other text.

from typing import List

def has_close_elements(numbers: List[float], threshold: float) -> bool:
    """
    """ Check if in given list of numbers, are any two numbers closer to each other than
    given threshold.
    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)
    False
    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
    True
    """
    """
"""
)
session_id, routing_target_model = client.model_select(
    messages=[{"role": "user", "content": new_prompt}],
    preference_id=preference_id,
)

print(f"Session ID: {session_id}")
print(f"Target Model: {routing_target_model}")

この例では、Not Diamond が Weave の自動トレースに対応している点も活用しています。結果は Weights & Biases UI で確認できます。

カスタムルーターを評価する

このセクションでは、カスタムルーターが単体で最も優れたモデルを上回るかどうかを測定する方法を説明します。カスタムルーターのトレーニングが完了したら、次のいずれかの方法でパフォーマンスを評価できます。

トレーニングに使用したプロンプトを送信して、サンプル内パフォーマンスを評価する。
新しいプロンプトまたはホールドアウトしたプロンプトを送信して、サンプル外パフォーマンスを評価する。

次の例では、テストセットをカスタムルーターに送信してパフォーマンスを評価します。

from weave.integrations.notdiamond.custom_router import evaluate_router

eval_prompt_column = "prompt"
eval_response_column = "actual"

best_provider_model, nd_model = evaluate_router(
    model_datasets=model_test,
    prompt_column=eval_prompt_column,
    response_column=eval_response_column,
    api_key=api_key,
    preference_id=preference_id,
)
python
@weave.op()
def is_correct(score: int, output: dict) -> dict:
    # モデルの応答はすでに取得済みなので、score をそのまま流用する簡易的な方法をとっています
    return {"correct": score}

best_provider_eval = weave.Evaluation(
    dataset=best_provider_model.model_results.to_dict(orient="records"),
    scorers=[is_correct],
)
await best_provider_eval.evaluate(best_provider_model)

nd_eval = weave.Evaluation(
    dataset=nd_model.model_results.to_dict(orient="records"), scorers=[is_correct]
)
await nd_eval.evaluate(nd_model)

この例では、Not Diamond の”メタモデル”が、プロンプトを複数の異なるモデルにルーティングします。 Weave を使ってカスタムルーターをトレーニングすると、評価も実行され、その結果が Weights & Biases UI にアップロードされます。カスタムルーターのプロセスが完了したら、Weights & Biases UI で結果を確認できます。 UI では、Not Diamond の”メタモデル”が、プロンプトに正確に回答できる可能性がより高い他のモデルへプロンプトをルーティングすることで、単体で最も高いパフォーマンスを示すモデルを上回っていることを確認できます。

​プロンプトのルーティング

​カスタムルーティング

​トレーニングデータを設定する

​カスタムルーターをトレーニングする

​カスタムルーターを評価する

プロンプトのルーティング

カスタムルーティング

トレーニングデータを設定する

カスタムルーターをトレーニングする

カスタムルーターを評価する